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
ActionSet
ActionSet
class max.experimental.sharding.ActionSet(axis_assignments, layouts, mesh, extras=(), result_shape=None, finalize=None)
Bases: object
A rule’s menu of per-axis sharding options for one op call.
Shape-aware (it depends on operand layouts) but cost-blind: it lists what is possible, not what is cheapest. The dispatcher picks one entry per mesh axis.
-
Parameters:
-
- axis_assignments (tuple[AxisAssignment, ...])
- layouts (tuple[TensorLayout, ...])
- mesh (DeviceMesh)
- extras (tuple[Any, ...])
- result_shape (Sequence[Dim] | None)
- finalize (Callable[[Action], Action] | None)
axis_assignments
axis_assignments: tuple[AxisAssignment, ...]
Per-axis rows, each pickable independently per mesh axis. The
last entry is the universal (R,…,R) -> R fallback.
extras
Non-tensor positional args appended after tensor-input mappings
in the picked Action’s inputs. A bare value is
treated as uniform across ranks; wrap in PerShard to vary
per rank.
finalize
Optional post-pick transform applied as finalize(action). Used
by rules that need the picked placement to compute per-rank metadata
or repack inputs into a user-facing container. Rules pre-bind any
per-op context by closing over it, so no separate context field is
needed.
layouts
layouts: tuple[TensorLayout, ...]
Per-tensor input layouts this menu was built for.
mesh
mesh: DeviceMesh
The mesh both the picker and planner work over.
result_shape
Output shape constraint, set by reshape-style rules. Lets the
feasibility check reject output Sharded rows whose result
dim is too small.