For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /get-started.md).

Python function

match_operand_placement

match_operand_placement()

max.experimental.sharding.match_operand_placement(graph_op, name, *, output_count=1)

source

Returns a sharding rule whose results take operand name’s placement.

The rule never moves an operand. It suits an op whose results are laid out like one of its inputs, such as per-row values computed from a row-sharded tensor.

Parameters:

• graph_op (Callable[[...], Any]) – The op the rule is for, read for its parameter names.
• name (str) – The parameter whose placement the results take.
• output_count (int) – The number of results. Defaults to 1.

Returns:

A rule that keeps every operand where it is and places each result as name is placed.

Raises:

ValueError – If graph_op has no parameter name.

Return type:

Callable[[…], list[AxisAssignment]]