SPMD Sharding
Logical axis partitioning, mesh mapping, and SPMD distribution rules.
-
taktiny.utils.spmd.with_logical_partitioning(initializer, axis_names=None, partition_spec=None)[source]
Wraps an initializer to instantly shard its output across the active JAX mesh
using logical axis rules, without allocating on a single device first.
- Parameters:
initializer (Any) – A standard JAX initializer function (key, shape, dtype).
axis_names (Sequence[str | None] | None) – The logical axis names for the parameter.
partition_spec (Any) – Optional explicit partition spec to use if axis_names is omitted or unmapped.
- Return type:
Any
- Returns:
A wrapped initializer function that automatically applies JIT and out_shardings
if a global JAX mesh is active.
-
taktiny.utils.spmd.set_logical_axis_rules(rules)[source]
- Return type:
None
- Parameters:
rules (LogicalRules)
-
taktiny.utils.spmd.get_logical_axis_rules()[source]
- Return type:
Sequence[tuple[str, str | tuple[str, ...] | None]]
-
taktiny.utils.spmd.map_logical_axis_names(map_names)[source]
Context manager and decorator for mapping logical axis names to mesh axis names.
Can also be called as a function without a ‘with’ block to set rules globally.
- Parameters:
map_names (Mapping[str, MeshAxisName] | Sequence[tuple[str, MeshAxisName]] | None)
- Return type:
None
-
taktiny.utils.spmd.logical_to_mesh_axes(array_dim_names, rules=None)[source]
- Return type:
P | None
- Parameters:
-