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
FeasibilityContext
FeasibilityContext
class max.experimental.sharding.FeasibilityContext(layouts, mesh, mesh_axis_idx, group_size, result_shape=None, input_placements=None, output_usage=None)
Bases: object
The state one feasibility check needs about one mesh axis.
Feasibility is decided per mesh axis: given the actual per-input
placements along that axis and the candidate AxisAssignment,
does the candidate produce non-empty shards on every tensor input,
and does its output dim fit the user-requested result shape (if
any)? This context carries everything the helpers in
action_is_feasible() consult.
-
Parameters:
group_size
group_size: int
mesh.mesh_shape[mesh_axis_idx]. Hoisted so the helpers don’t keep
re-indexing.
input_placements
input_placements: Sequence[tuple[Placement, ...]] | None = None
Current per-input placements (one tuple per tensor input). None
falls back to reading layouts[i].mapping.
layouts
layouts: tuple[TensorLayout, ...]
Per-tensor input layouts. Indexed positionally; one entry per tensor argument the rule received.
mesh
mesh: DeviceMesh
The device mesh; only the per-axis size is read.
mesh_axis_idx
mesh_axis_idx: int
The mesh axis under evaluation. Indexes mesh.mesh_shape and the
per-axis tuple of every input’s placement.
output_usage
Per-tensor-axis cumulative shard factor across mesh axes already picked in a multi-axis search. Lets the output-dim check stay consistent when more than one mesh axis shards the same tensor axis.
result_shape
Output shape constraint set by reshape-style rules (rules where the
output shape is a rule argument, not derived). Used to reject
Sharded output rows whose target dim is smaller than the
cumulative shard factor.
Was this page helpful?
Thank you! We'll create more content like this.
Thank you for helping us improve!