Sharding Rules

Named-axis sharding rules and partition-spec helpers for mapping logical axes onto a device mesh.

See Also

• Distributed Overview - All distributed tools
• Device Mesh - Mesh construction
• Data Parallel - Data parallelism
• Distributed Training Guide

---

datarax.distributed.sharding

NNX SPMD sharding utilities for Datarax.

This module provides utilities for configuring sharding strategies in JAX SPMD training, following upstream Flax NNX patterns.

The core abstraction is MeshRules, which maps logical dimension names (data, embed, mlp, heads) to physical mesh axis names. Factory functions create standard configurations for data-parallel and FSDP training.

Example

mesh = jax.make_mesh((4, 2), ("data", "model")) rules = fsdp_rules(data_axis="data", model_axis="model") spec = partition_spec_for_names(rules, "data", "embed")

=> PartitionSpec("data", "model")

logger module-attribute

logger = logging.getLogger(__name__)

MeshRules dataclass

MeshRules(*, data: str | None = None, embed: str | None = None, mlp: str | None = None, heads: str | None = None)

Logical-to-physical mesh axis mapping for SPMD sharding.

Maps logical dimension names to physical mesh axis names. Following the pattern from Flax NNX FSDP examples.

Attributes:

Name	Type	Description
data	str | None	Mesh axis for the batch/data dimension.
embed	str | None	Mesh axis for embedding dimensions.
mlp	str | None	Mesh axis for MLP hidden dimensions.
heads	str | None	Mesh axis for attention head dimensions.

data class-attribute instance-attribute

data: str | None = None

embed class-attribute instance-attribute

embed: str | None = None

mlp class-attribute instance-attribute

mlp: str | None = None

heads class-attribute instance-attribute

heads: str | None = None

data_parallel_rules

data_parallel_rules(data_axis: str = 'data') -> MeshRules

Create MeshRules for pure data-parallel training.

Shards only the batch dimension across devices.

Parameters:

Name	Type	Description	Default
data_axis	str	Name of the mesh axis for data parallelism.	'data'

Returns:

Type	Description
MeshRules	MeshRules with only the data axis mapped.

fsdp_rules

fsdp_rules(data_axis: str = 'data', model_axis: str = 'model') -> MeshRules

Create MeshRules for Fully Sharded Data Parallel (FSDP) training.

Shards the batch dimension across the data axis and model weight dimensions across the model axis.

Parameters:

Name	Type	Description	Default
data_axis	str	Name of the mesh axis for data parallelism.	'data'
model_axis	str	Name of the mesh axis for model parallelism.	'model'

Returns:

Type	Description
MeshRules	MeshRules with data and model axes mapped.

create_named_sharding

create_named_sharding(mesh: Mesh, *axis_names: str | None) -> NamedSharding

Create a NamedSharding from a mesh and axis names.

Parameters:

Name	Type	Description	Default
mesh	Mesh	The device mesh.	required
*axis_names	str | None	Axis names for each dimension. Use None for replicated dimensions. If no axes are provided, creates a fully replicated sharding.	()

Returns:

Type	Description
NamedSharding	A JAX NamedSharding object.

partition_spec_for_names

partition_spec_for_names(rules: MeshRules, *logical_names: str) -> PartitionSpec

Apply MeshRules to create a PartitionSpec for a tensor.

Maps logical dimension names through the rules to produce physical mesh axis assignments.

Parameters:

Name	Type	Description	Default
rules	MeshRules	The MeshRules mapping to apply.	required
*logical_names	str	Logical dimension names for each tensor axis.	()

Returns:

Type	Description
PartitionSpec	A PartitionSpec with physical mesh axis assignments.