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

auto_reshard

auto_reshard()

max.experimental.sharding.auto_reshard(allowed_transitions=None, *, mode=None)

source

Sets how ops in this block may change the placement of their inputs.

On each mesh axis the picker takes the cheapest row whose transitions are all in allowed_transitions, and mode decides what happens when that row moves an input. Outside any block the picker allows DEFAULT_TRANSITIONS and moves inputs silently; an argument left out keeps the enclosing block’s setting. Use "warn" to find the moves an op makes, and "raise" to require each one to be written as an explicit transfer_to. Moving an input from another mesh onto the op’s mesh counts as a move, whatever its placement there. Also works as a decorator. For example:

with auto_reshard(mode="warn"):    # reshard, and say so each time
    ...
with auto_reshard(mode="raise"):   # require explicit transfer_to
    ...
with auto_reshard(ALL_TRANSITIONS):      # sequence parallelism too
    ...
with auto_reshard({(Partial, Replicated)}):
    ...

Parameters:

• allowed_transitions (Iterable[tuple[type[Placement], type[Placement]]] | None) – The (from, to) placement-type pairs the picker may plan with. A row needing any other transition is not allowed; if no row is allowed on some mesh axis, the op raises.
• mode (Literal['silent', 'warn', 'raise'] | None) – What happens when a chosen row moves an input. "silent" moves it, "warn" moves it and emits a UserWarning naming the move and the transfer_to that would make it explicit, "raise" refuses with a ShardingError saying the same.

Raises:

• TypeError – If an allowed transition is not a pair of placement types.
• ValueError – If mode is not one of the three above.

Return type:

Generator[None]