NNX State Sharding Specifications

Why this matters

A model can contain many arrays, each needing a different device layout. A matching tree of partition specifications lets you describe those layouts alongside the model state. This exercise constructs that specification tree for one NNX Linear layer.

The required function returns the number of specifications. It does not create a mesh, move arrays, or run distributed computation. A count of two is a useful structure check, but cannot prove that the specifications are correct or that any array has been sharded.

Separate graph structure from variable state

graphdef, state = nnx.split(model) separates the graph description from the model's variable state. The state preserves variable types and metadata; JAX pytree traversal exposes the underlying array leaves. For a default nnx.Linear(in_dim, features), those arrays are:

• A kernel with shape (in_dim, features).
• A bias with shape (features,).

The graph definition records how to reconstruct the module. After transforming the state, use nnx.merge(graphdef, new_state) when you need a callable model again. Splitting is useful for this explicit state workflow; it is not the only way to work with NNX transforms.

Assign one specification to each array

For this layer, choose P(None, "model") for the kernel and P("model") for the bias. The kernel's output columns and the corresponding bias entries then use the same model mesh axis. The kernel's input dimension stays whole in each shard.

In this particular state tree, rank identifies the two arrays: rank two means kernel, rank one means bias. That rule is not a general sharding policy. A larger model can contain embeddings, normalization parameters, convolution kernels, and other arrays whose layouts depend on their role, not just their rank.

Apply the rule with jax.tree_util.tree_map. In this workflow, PartitionSpec objects are pytree leaves, so the resulting tree has two leaves even though each specification contains axis entries. Do not count the strings or None entries inside a specification.

A complete example

The following example builds the specification tree, applies it to one available device, merges the state, and checks the forward result:

import jax
import jax.numpy as jnp
import numpy as np
from flax import nnx
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P

model = nnx.Linear(3, 4, rngs=nnx.Rngs(0))
graphdef, state = nnx.split(model)

def pick_spec(array):
    return P(None, "model") if array.ndim == 2 else P("model")

specs = jax.tree_util.tree_map(pick_spec, state)
assert len(jax.tree_util.tree_leaves(specs)) == 2
assert jax.tree_util.tree_structure(specs) == jax.tree_util.tree_structure(state)

mesh = Mesh(np.array(jax.devices()[:1]), ("model",))
placed = jax.tree_util.tree_map(
    lambda array, spec: jax.device_put(array, NamedSharding(mesh, spec)),
    state, specs,
)
restored = nnx.merge(graphdef, placed)
x = jnp.ones((2, 3))
np.testing.assert_allclose(restored(x), model(x), rtol=1e-5, atol=1e-6)
assert restored.kernel[...].sharding.spec == P(None, "model")
assert restored.bias[...].sharding.spec == P("model")

With one device, each shard still contains the full array. With multiple model-axis devices and an evenly divisible output width, these specifications divide the kernel columns and bias entries. Actual placement requires a mesh and a sharding operation; building the specification tree alone changes no array's location.

Matching kernel and bias layouts is convenient, but a replicated bias does not inherently produce incorrect arithmetic. Whether an operation accepts a combination, needs an explicit annotation, or performs resharding depends on its sharding mode and compilation. Layout affects execution and communication as well as local shapes.