Skip to content

← Flax NNX step 83 of 100

Medium Primitives

NNX Initialization State and Mesh Configuration

Jump to the task

Why this matters

Before distributing a model, you need to know which arrays it owns and how you want to place them. These are separate questions. An NNX state tree describes model variables; a device mesh names a layout of devices; a sharding specification connects array dimensions to that layout.

This exercise covers the first step: initialize a linear layer and inspect its state. It also records a proposed mesh size. It does not create a mesh, distribute arrays, simulate communication, or measure memory use. The historical function name mesh_init_simulate is retained, but its output is only a configuration summary.

Inspect the initialized layer

For nnx.Linear(in_dim, features) with its default enabled bias, the parameter arrays are:

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

graphdef, state = nnx.split(model) separates the model's graph description from its variable state. For this layer, traversing state with jax.tree_util.tree_leaves gives two array leaves. Two leaves does not mean two scalar parameters: the scalar count is in_dim * features + features.

This complete example has two leaves containing 15 scalar values:

import jax
import jax.numpy as jnp
from flax import nnx

model = nnx.Linear(4, 3, rngs=nnx.Rngs(0))
graphdef, state = nnx.split(model)
leaves = jax.tree_util.tree_leaves(state)
assert len(leaves) == 2
assert sum(a.size for a in leaves) == 15
restored = nnx.merge(graphdef, state)
assert restored(jnp.ones((2, 4))).shape == (2, 3)

Count the state of the model you construct. Other modules may have additional state, and disabling this layer's bias changes its leaf count. A leaf count alone cannot verify parameter values or placement.

What a mesh adds

A mesh arranges JAX devices along named axes. For example, an axis named "model" with size four contains four devices. It is still one axis. This exercise's mesh_axes argument means the proposed size of that one axis, not the number of axes.

Mesh names do not select a parallelization algorithm. A model axis can be unused by an array's partition specification, in which case that array is replicated over that axis. A two-dimensional mesh does not automatically implement data or tensor parallelism.

NamedSharding(mesh, spec) explicitly combines a mesh and a PartitionSpec. Here is a runnable one-device example, separate from the required entrypoint:

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

mesh = Mesh(np.array(jax.devices()[:1]), ("model",))
layout = NamedSharding(mesh, P(None, "model"))
array = jax.device_put(jnp.arange(12.).reshape(3, 4), layout)
assert array.shape == (3, 4)
assert array.sharding.spec == P(None, "model")
assert len(array.addressable_shards) == 1

The specification assigns the second array dimension to the model axis. With only one device, the local shard is still the whole array. This shows the layout API, not a memory saving or distributed run.

Large-model initialization can use abstract shape evaluation and compiled initialization with specified output shardings to avoid first creating an eager full parameter array on one device. The exact method depends on the JAX/NNX sharding mode and version. A mesh context alone does not specify how every array dimension is divided, and output placement alone is not proof of peak-memory behavior.

Problem

Implement mesh_init_simulate(seed, x, features, mesh_axes):

  1. Cast seed, features and mesh_axes to Python integers. Assume positive feature and proposed mesh sizes.
  2. Use x.shape[-1] as in_dim. Build nnx.Linear(in_dim, features, rngs=nnx.Rngs(seed)), with default initializers and an enabled bias.
  3. Split the model with nnx.split(model) and count the array leaves of its state with jax.tree_util.tree_leaves.
  4. Return a floating-point array of shape (3,) containing [features, mesh_axes, n_state_leaves], in that order.

Input x has shape (N, in_dim). Only its last dimension is used; do not run the layer on x. Its values and batch size do not affect the report. The seed changes initialized kernel values but not their number or shapes. The proposed mesh size is reported unchanged; it need not equal the available device count or divide the output width, because this entrypoint does not apply a partition specification.

The diagram below is a separate two-axis sharding illustration. It starts fully replicated. Changing its selectors shows how an array's partition specification, rather than mesh names alone, determines which dimensions are split. It does not visualize the arrays created by your function.

Loading visualization…