For the complete documentation index, see llms.txt. Markdown versions of all pages are available by appending .md to any URL (e.g. /max/get-started.md).
Python class
ActionSet
ActionSet
class max.experimental.sharding.ActionSet(axis_assignments, layouts, mesh, extras=(), result_shape=None, finalize=None, finalize_ctx=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, Any], Action] | None)
- finalize_ctx (Any)
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, finalize_ctx). Used by rules that need the
picked placement to compute per-rank metadata or repack inputs
into a user-facing container.
finalize_ctx
finalize_ctx: Any = None
Opaque per-op context passed to finalize.
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.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!