Skip to content

← Flax NNX step 81 of 100

Medium Primitives

NNX PartitionSpec Layout

Jump to the task

Why this matters

A sharding layout tells JAX which parts of an array belong on which devices. Before applying a layout, you can work out the shape each device will hold. This exercise asks for that shape for one specific layout: split a linear layer's output columns over a model axis.

The function only returns two dimensions. It does not build an NNX layer, allocate a kernel, create a mesh, or distribute computation. The calculation applies to an NNX Linear kernel because its shape is (input_features, output_features).

Read the specification left to right

A PartitionSpec describes array dimensions using names from a device mesh. For a matrix, P(None, "model") means:

  • Dimension 0 is not split, so each shard keeps every input row.
  • Dimension 1 is divided over the mesh axis named "model".

If the matrix has shape (8, 16) and the model axis has size four, the four distinct shards each have shape (8, 4). They contain columns 0 through 3, 4 through 7, 8 through 11, and 12 through 15. The complete kernel is not replicated on each model-axis device.

Here is a complete NumPy illustration of those column slices. It demonstrates the contents without requiring multiple JAX devices:

import numpy as np

kernel = np.arange(8 * 16).reshape(8, 16)
shards = np.split(kernel, 4, axis=1)
assert all(shard.shape == (8, 4) for shard in shards)
np.testing.assert_array_equal(shards[1], kernel[:, 4:8])
np.testing.assert_array_equal(np.concatenate(shards, axis=1), kernel)

Array dimensions and mesh axes are different

For the fully specified layouts discussed here:

Spec entry Effect on that array dimension
None Keep its full length in each shard
"model" Divide its length by the model-axis size
("data", "model") Divide its length by the product of those axis sizes

Unused mesh axes replicate the resulting shards. For example, on a mesh of shape (2, 4) named ("data", "model"), P(None, "model") creates four distinct column shards, each held by two devices. P("data", "model") instead splits both matrix dimensions and produces eight distinct shards. These are different layouts.

A mesh axis can be assigned to at most one array dimension in these specifications. The assigned dimensions must divide evenly by their mesh-axis sizes for the equal-sized shards used in this exercise. Axis names identify devices, not the semantic meaning of the array.

For example, P("data", None) on an input batch shards examples; the same spec on a weight matrix shards its rows. A batch layout alone is not fully sharded data parallel training. That also requires decisions about parameters, gradients, optimizer state and execution.

Short specifications are valid

Unspecified trailing array dimensions are treated as unpartitioned. For a matrix, P("model") has the same layout as P("model", None): it splits rows. It does not mean P(None, "model"), which splits columns. P() leaves all dimensions unpartitioned. A specification with too many dimensions is invalid; it does not broadcast to make the ranks match.

To place an array explicitly, combine the spec with an actual mesh using NamedSharding(mesh, spec). This runnable example uses one available device, so its local shard remains the full matrix:

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.addressable_shards[0].data.shape == (3, 4)

This demonstrates placement syntax. It does not claim distributed speedups, memory savings, or a particular parallel linear algorithm.

Problem

Implement partition_spec_layout(num_devices, kernel_in, kernel_out). Cast all three scalar inputs to Python integers. Assume positive values and kernel_out divisible by num_devices.

Here num_devices is the size of the one mesh axis named model. Return a floating-point array of shape (2,) containing:

  1. kernel_in, since the input-feature dimension is not split.
  2. kernel_out // num_devices, the output columns in one shard.

The input numbers describe a proposed layout and do not depend on the devices available to run your solution. Do not divide the input dimension, return the total number of elements, or create devices. With one proposed device, the shard shape equals the full shape.

The diagram below starts with an 8-by-12 array on a 2-by-4 mesh and the column-sharding layout used here. Its unused data axis creates two copies of each 8-by-3 column shard. Changing the selectors lets you compare that with layouts that also split rows.

Loading visualization…