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 class
AxisAssignment
AxisAssignment
class max.experimental.sharding.AxisAssignment(needed_inputs, outputs)
Bases: NamedTuple
One placement choice that a sharding rule offers for its op.
A sharding rule is a function with the same parameters as its op. It
receives the op’s arguments with defaults filled in, so every parameter of
the op is present, and each tensor, at any depth, replaced by its
TensorLayout. It returns a list of rows. Each row gives a
placement for every tensor operand, in argument order, and for every
result. Some rows of the matmul rule:
from max.experimental.sharding import (
AxisAssignment,
Partial,
Replicated,
Sharded,
TensorLayout,
)
def matmul_rule(
lhs: TensorLayout, rhs: TensorLayout
) -> list[AxisAssignment]:
return [
AxisAssignment((Replicated(), Replicated()), (Replicated(),)),
# Column parallel: each device holds some of the columns.
AxisAssignment((Replicated(), Sharded(1)), (Sharded(1),)),
# Row parallel: each device holds part of the sum.
AxisAssignment((Sharded(1), Sharded(0)), (Partial(),)),
]On each mesh axis, the op discards every row that needs a transition
outside the allowed transition set, then takes the cheapest of the rows
that remain; ties go to the row that moves fewer inputs. It moves its
operands to that row’s placements and gives its results the row’s
placements. The op never moves an operand whose own mesh holds a single
device, so every row must name Replicated for it, and a row that
names anything else is infeasible. When no row survives on a mesh axis,
the op raises ShardingError in every mode. A rule never changes
the op’s other arguments.
Create new instance of AxisAssignment(needed_inputs, outputs)
needed_inputs
The placement of each tensor operand, in argument order.
outputs
The placement of each result, in result order.