Skip to content

← Flax step 89 of 100

Medium Primitives

PartitionSpec Layout

Jump to the task

Calculate the shape of a kernel shard

A Dense kernel has shape (kernel_in, kernel_out). Suppose its output columns are divided evenly across a device-mesh axis named model, while every shard keeps all input rows. Compute the shape of one shard.

This layout is described by PartitionSpec(None, "model"), often written P(None, "model"). The function performs shape arithmetic only: it does not create a mesh, allocate a kernel or move data.

Read the specification by array dimension

A PartitionSpec associates array dimensions with named mesh axes. For this two-dimensional kernel:

Kernel dimension Entry Effect
Input rows None Keep all kernel_in rows in each shard.
Output columns "model" Divide kernel_out by the size of the model mesh axis.

None means this array dimension is not partitioned. It does not mean the whole array is copied onto every device. Replication depends on which mesh axes the complete specification leaves unused.

In more general specifications, a tuple such as ("data", "model") in one entry partitions that array dimension across the product of those mesh-axis sizes. This exercise uses only the two entries above.

Distinguish mesh-axis size from total device count

Imagine eight devices arranged as data=2, model=4. For a kernel of shape (8, 12) and P(None, "model"):

  • The four model positions hold four different column blocks.
  • Each block has shape (8, 3) because 12 / 4 = 3.
  • The unused data axis repeats each block twice.

Here the function's num_devices argument is 4, the size of the model axis, not 8, the total devices in the mesh. The full logical kernel remains (8, 12).

Array dimensions have meanings supplied by the surrounding computation. For a Dense kernel they are input and output features; its first axis is not a batch axis. A specification alone does not identify an entire training strategy or determine every collective operation it needs.

See the column blocks without multiple devices

This local array example shows exactly what the four distinct blocks contain. Replication on the unused data axis would copy these blocks; it would not cut them into smaller pieces.

import jax.numpy as jnp

kernel = jnp.arange(8 * 12).reshape(8, 12)
blocks = jnp.split(kernel, 4, axis=1)
assert len(blocks) == 4
assert all(block.shape == (8, 3) for block in blocks)
assert jnp.array_equal(blocks[0], kernel[:, :3])
assert jnp.array_equal(blocks[3], kernel[:, 9:])
assert jnp.array_equal(jnp.concatenate(blocks, axis=1), kernel)

To place such an array on actual devices, combine a compatible mesh and specification in jax.sharding.NamedSharding, then use that sharding with jax.device_put. That requires enough devices for the mesh. The arithmetic in this exercise has no hardware-count requirement.

Problem

Implement partition_spec_layout(num_devices, kernel_in, kernel_out):

  1. Convert all three values to Python integers.
  2. Keep the row size equal to kernel_in.
  3. Compute the column size as kernel_out // num_devices.
  4. Return those two sizes as a one-dimensional float array of shape (2,), in row-then-column order.

All values are positive integers, possibly supplied as integer-valued Python floats. kernel_out is exactly divisible by num_devices; no rounding, padding or uneven shards are part of this task. With a model-axis size of one, the shard has the full kernel shape.

Keep all three arguments static if JIT-compiling this implementation, since it converts them to Python integers. Return dimensions, not array contents, a PartitionSpec object or the number of replicas.

The diagram starts with the (8, 12) kernel on the data=2, model=4 mesh described above. Its controls let you compare other layouts; the function always uses the initial P(None, "model") layout.

Loading visualization…