Simulate a Sharded Dense Forward Pass

Problem

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

1. Convert seed, num_shards, and features to Python integers.
2. Create nn.Dense(features) with its default initialization and bias.
3. Set chunk_size = x.shape[0] // num_shards.
4. Initialize the layer once on the first chunk, using jax.random.PRNGKey(seed) without splitting the key.
5. Apply that same initialized layer to each consecutive chunk of x.
6. Concatenate the outputs along axis 0 and return the flattened array in row-major order.

Inputs:

• seed: an integer-valued random seed.
• x: a float array of shape (N, in_dim).
• num_shards: a positive integer-valued chunk count that divides N.
• features: a positive integer-valued output width.

The scalar inputs may arrive as floats; convert them with int. The batch and input width are positive, and each chunk has at least one row. Return a one-dimensional array of length N * features. With num_shards=1, apply the layer to the entire batch.

Check the shapes

For x.shape == (6, 4), num_shards == 3, and features == 2:

Use concatenate, not stack: stacking introduces a new chunk axis. You can slice with x[i * chunk_size:(i + 1) * chunk_size] or use jax.lax.dynamic_slice_in_dim. In this Python loop, the slice bounds are known while tracing. If you wrap the whole function in jax.jit, keep seed, num_shards, and features static.

Dense initializes its kernel from the last input dimension. Initializing on the first chunk establishes the parameter shapes; it does not make the layer specific to that batch size. Reuse the complete variables tree returned by init in every call to apply.

How this relates to real shard_map

jax.shard_map runs a function on device-local blocks. Its input specs describe those blocks; its output specs describe how the local results form a global array. This complete example uses the process's local devices and replicates the Dense parameters across them:

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

devices = np.array(jax.local_devices())
mesh = Mesh(devices, axis_names=("data",))
x = jnp.arange(8 * len(devices), dtype=jnp.float32).reshape(-1, 4)
model = nn.Dense(2)
variables = model.init(jax.random.PRNGKey(0), x[:2])

def local_apply(variables, local_x):
    return model.apply(variables, local_x)

sharded_apply = jax.shard_map(
    local_apply,
    mesh=mesh,
    in_specs=(P(), P("data", None)),
    out_specs=P("data", None),
)
y = sharded_apply(variables, x)
np.testing.assert_allclose(y, model.apply(variables, x), rtol=1e-5, atol=1e-5)
assert y.shape == (2 * len(devices), 2)

P() replicates each parameter array. P("data", None) splits rows across the data axis while keeping all features in each local block. The output remains a distributed array; reconstructing its logical shape does not mean gathering every value onto one device.

No collective operation is needed inside this particular forward pass because rows are independent. A batch-wide reduction or training step needs additional coordination. For example, normalizing each chunk by its own mean generally differs from normalizing by the full batch mean. The Python simulation does not model communication or guarantee parallel execution, and its output tests cannot prove that multiple devices ran. See the shard_map reference for the actual API.

The diagram illustrates row splitting: eight rows on two data devices, each holding a (4, 12) block. The model axis has size one, so features are not split.