Sharding across devices

Sharding places different parts of an array on different devices. Write operations on the whole array; JAX runs them across its shards and handles communication between devices.

Taktiny uses JAX’s Mesh, PartitionSpec, and NamedSharding. Logical axis names let a model describe its dimensions independently of a device layout.

Mesh axes and array dimensions

Object	Meaning	Example
Mesh	An arrangement of devices with named axes.	A 2 × 2 mesh named ("data", "model").
PartitionSpec	Which mesh axes partition each array dimension.	P(None, "model") partitions the second dimension.
NamedSharding	A spec attached to a particular mesh.	NamedSharding(mesh, P("data", None)).
axis_names	Logical names for a parameter’s dimensions.	("input", "output").

In a spec, None leaves that array dimension unpartitioned. Mesh axes unused by the entire spec replicate the array. For example, P(None, "model") partitions a matrix’s columns along model and replicates it along data. P() requests full replication.

Each entry in a spec describes one array dimension. Its size must be divisible by the product of the sizes of its assigned mesh axes.

A complete four-device example

Try the following example on four CPU devices. Save the Python block as sharding_example.py and run:

JAX_PLATFORMS=cpu XLA_FLAGS='--xla_force_host_platform_device_count=4' uv run python sharding_example.py

Set the environment before JAX initializes its backend. The flag exposes four logical CPU devices for testing.

import jax
import jax.numpy as jnp
import numpy as np
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P
from taktiny import nn

assert len(jax.devices()) == 4, "Run this example with four JAX devices"
mesh = Mesh(np.array(jax.devices()).reshape(2, 2), ("data", "model"))

with jax.set_mesh(mesh):
    layer = nn.Linear(
        8, 12,
        partition_spec=P(None, "model"),
        rngs=nn.Rngs(0),
    )
    x = jax.device_put(
        np.ones((4, 8), dtype=np.float32),
        NamedSharding(mesh, P("data", None)),
    )
    forward = jax.jit(
        lambda model, inputs: model(inputs),
        out_shardings=NamedSharding(mesh, P("data", "model")),
    )
    y = forward(layer, x)

assert layer.kernel.value.shape == (8, 12)
assert layer.kernel.value.sharding.spec == P(None, "model")
assert layer.bias.value.sharding.spec == P("model")
assert y.shape == (4, 12)
assert y.sharding.spec == P("data", "model")

The kernel’s output dimension and the bias are partitioned across model. Inputs are partitioned across data. The requested output layout uses both: each device holds part of the batch and part of the output features.

The remaining Python blocks build on this example’s imports and mesh.

Name dimensions, then choose their placement

Use map_logical_axis_names to map model-level names to mesh axes:

from taktiny.utils.spmd import map_logical_axis_names

with jax.set_mesh(mesh), map_logical_axis_names({
    "input": None,
    "output": "model",
}):
    logical_layer = nn.Linear(
        8, 12,
        axis_names=("input", "output"),
        rngs=nn.Rngs(1),
    )

assert logical_layer.kernel.axis_names == ("input", "output")
assert logical_layer.kernel.partition_spec == P(None, "model")
assert logical_layer.kernel.value.sharding.spec == P(None, "model")

Taktiny preserves the original axis_names and stores the resolved mesh spec separately. For Linear, the bias uses the trailing output axis names and spec entries.

The rules apply when the initializer or parameter is created. The with block restores the previous rules on exit. To move an existing array to a new layout, use jax.device_put with the desired sharding.

When both arguments are supplied

axis_names takes precedence over partition_spec: the resolved logical names replace the entire explicit spec, including unmapped dimensions:

with jax.set_mesh(mesh), map_logical_axis_names({"output": "model"}):
    overridden = nn.Linear(
        8, 12,
        axis_names=("input", "output"),
        partition_spec=P("data", None),
        rngs=nn.Rngs(2),
    )

assert overridden.kernel.partition_spec == P(None, "model")

Here input is unmapped, so it resolves to None. If no logical names match any rules, every dimension resolves to None and the parameter is replicated on the active mesh.

For an explicit layout, pass partition_spec on its own. To replicate a parameter while keeping logical names, map those names to None when creating it.

Rule order and scope

Preview a mapping with logical_to_mesh_axes:

from taktiny.utils.spmd import logical_to_mesh_axes

assert logical_to_mesh_axes(
    ("batch", "feature"),
    rules=(("batch", "data"), ("feature", "model")),
) == P("data", "model")

assert logical_to_mesh_axes(
    ("input", "output"),
    rules=(("output", "model"), ("input", "model")),
) == P(None, "model")

Rules are considered in order. Each mesh axis can partition one dimension of an array, so the first available assignment wins. In the second example, output claims model, leaving input unpartitioned. Give each named array dimension a distinct logical name.

A logical name can map to a tuple of mesh axes. For example, {"feature": ("data", "model")} partitions one dimension across both axes; its size must be divisible by four on this mesh.

Prefer with map_logical_axis_names(...) for temporary rules. Nested blocks prepend their rules, giving them priority over outer rules, and restore the previous rules on exit.

For longer-lived configuration, set_logical_axis_rules(rules) replaces the current rule sequence; get_logical_axis_rules() reads it. Set rules in each thread or process that constructs parameters.

Calling map_logical_axis_names(...) without with prepends rules immediately. Use set_logical_axis_rules(...) when you want to replace the full sequence.

When initialization is sharded

Layers such as Linear wrap their initializers with with_logical_partitioning when sharding arguments are supplied. With an active jax.set_mesh(mesh) context and a resolved spec, the wrapper uses jax.jit(..., out_shardings=spec). The initialized parameter already has its requested layout. The output layout describes the returned array; the initializer may still use temporary storage while it runs.

Construct parameters inside the mesh context for placement at initialization. Outside that context, the initializer returns an ordinary array while the parameter can still retain its sharding metadata. Check parameter.value.sharding to inspect the actual layout.

You can wrap a custom initializer directly:

from taktiny.utils.spmd import with_logical_partitioning

initialize = with_logical_partitioning(
    jax.nn.initializers.normal(),
    axis_names=("input", "output"),
)
with jax.set_mesh(mesh), map_logical_axis_names({"output": "model"}):
    weight = initialize(jax.random.key(3), (8, 12), jnp.float32)

assert weight.sharding.spec == P(None, "model")

Use the usual initializer signature (key, shape, dtype). Logical rules are resolved when the wrapped initializer is called.

Match names to the parameter shape

A parameter needs one logical axis entry per array dimension. A feature shape with several dimensions therefore needs several entries:

with jax.set_mesh(mesh), map_logical_axis_names({"output": "model"}):
    nd_layer = nn.Linear(
        (2, 4), (3, 4),
        axis_names=("input_group", "input_feature", "output_group", "output"),
        rngs=nn.Rngs(4),
    )

assert nd_layer.kernel.value.shape == (2, 4, 3, 4)
assert nd_layer.kernel.partition_spec == P(None, None, None, "model")
assert nd_layer.bias.value.shape == (3, 4)
assert nd_layer.bias.partition_spec == P(None, "model")

Check each layer’s parameter layout when choosing names: convolution kernels also have spatial dimensions, and bilinear kernels have dimensions for both inputs before the output dimensions.

Inspect and place arrays

For an ordinary array-backed parameter, parameter.value.sharding reports its actual placement. array.addressable_shards exposes the shards accessible from the current process, including their global indices and local arrays. Replication can mean several devices hold the same global slice.

Use jax.device_put(array, NamedSharding(mesh, spec)) to explicitly place or reshard an existing array. For input batches, either use jax.device_put as in the first example or configure DataLoader with axis_names or partition_spec to place completed batches. Choose output layouts at the JIT boundary when your application needs a particular layout.

DataLoader(shard_index=..., shard_count=...) assigns source records to processes; its output-placement arguments control the resulting JAX arrays. These examples use one process with four devices. For multi-process execution, initialize JAX distributed execution and construct global arrays from the process-local data.

See the SPMD API reference for mapping helpers and the data guide for input-pipeline sharding.