Skip to content

← Flax NNX step 87 of 100

Hard Primitives

NNX shard_map Simulation

Jump to the task

Why this matters

A linear layer applies the same affine transformation to each batch row. You can divide those rows into groups, apply one shared layer to every group, and concatenate the results in their original order. This is the forward computation behind a simple data-parallel layout.

The required function simulates that computation with a local loop. It does not call shard_map, create a mesh, or execute groups on separate devices. The full model, input and output remain present. The result checks numerical behavior, not parallel speed or memory use.

Local blocks and the global array

JAX's shard_map expresses a function in terms of blocks assigned to devices. With a mesh axis named data and input specification P("data", None), a matrix's rows are split across that axis while each block keeps its full feature dimension.

If input x has shape (N, D) and the data axis has size S, the local block has shape (N // S, D), assuming exact division. This preserves array rank: mapping rows into blocks is different from passing one row at a time to a function.

Output specification P("data", None) assembles local output blocks along the global batch axis. It does not by itself gather the entire result onto every device. The result may remain a distributed array whose global shape is (N, features).

A real one-device example

shard_map also works with a one-device mesh. That does not divide the array into smaller pieces, but it makes a complete example runnable without requiring several accelerators:

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

mesh = Mesh(np.array(jax.devices()[:1]), ("data",))
x = jnp.array([[1., 2.], [3., 4.]])
kernel = jnp.array([[1., -1.], [2., 0.]])
bias = jnp.array([0.5, 1.])

def local_apply(local_x, full_kernel, full_bias):
    return local_x @ full_kernel + full_bias

apply = jax.shard_map(
    local_apply, mesh=mesh,
    in_specs=(P("data", None), P(), P()),
    out_specs=P("data", None),
)
result = apply(x, kernel, bias)
np.testing.assert_allclose(result, [[5.5, 0.], [11.5, -2.]])
np.testing.assert_allclose(result, x @ kernel + bias)

The empty specifications for kernel and bias replicate those arrays over the mesh. With more devices, this layout gives each device different input rows and the same parameters. This independent forward operation needs no cross-device reduction. Computations needing sums or gathered values across devices require the appropriate collective operations in their manual parallel implementation.

Simulate the row groups

Build one nnx.Linear(D, features, rngs=nnx.Rngs(seed)) with default initialization and enabled bias. For group i, select rows from i * chunk through (i + 1) * chunk, where chunk = N // num_shards. Apply the same model to each group without changing its parameters.

Each result has shape (chunk, features). Concatenating along axis 0 joins existing rows and gives (N, features). Stacking would add an extra axis, producing (num_shards, chunk, features) instead.

Initializing a layer inside the loop is unnecessary. Reusing the same seed could recreate the same weights, while advancing an RNG could produce different weights. The requirement here is simpler: create one model once, and use it for every slice.

Slicing and tracing

Ordinary slicing works when its boundaries are Python integers. You may also use jax.lax.dynamic_slice_in_dim with a dynamic start index, but its slice size must still be static for this JIT workflow. It does not make arbitrary traced sizes valid.

The seed, shard count and feature count determine construction or loop structure here. Treat them as static arguments when checking this entrypoint under JIT. Input array values can remain dynamic.

Problem

Implement shard_map_simulate(seed, x, num_shards, features):

  1. Cast seed, num_shards and features to Python integers. Assume positive dimensions and N divisible by the positive shard count.
  2. Create the Linear once, using x.shape[-1] as input width.
  3. Divide x into num_shards contiguous axis-0 slices of length N // num_shards. Apply the shared model to each slice.
  4. Concatenate results along axis 0 in original row order.
  5. Return out.reshape(-1), with length N * features.

The shard count describes loop groups and need not match available devices. One group uses the full batch; N groups use one example each. The result should match applying the same full Linear to x within floating-point tolerance.

This equivalence is specific to the row-independent layer here. Shard-local BatchNorm statistics or differing randomness in a more complex model could change the result. This exercise includes no optimizer, training update, or mutable forward state.

The diagram illustrates batch-row sharding on a separate two-axis mesh. Its unused model axis replicates each row block. It does not show actual placement by the required local simulation.

Loading visualization…