Skip to content

← JAX Performance & Distributed step 12 of 25

Medium Primitives

PartitionSpec Basics

Jump to the task

Why this matters

A mesh names groups of devices. A PartitionSpec maps array dimensions onto those mesh axes. For a matrix, PartitionSpec('data', None) splits rows across the data axis and leaves the columns unsplit in each row block.

None means no partitioning of that array dimension. Replication occurs across mesh axes unused by the specification. For example, on a mesh with data and model axes, this row-only specification repeats each row block across the model axis. It does not give every device the whole matrix.

Worked mini-example

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

mesh = Mesh(np.array(jax.devices()[:1]), ('data',))
sharding = NamedSharding(mesh, PartitionSpec('data', None))
x = jnp.arange(12, dtype=jnp.float32).reshape(4, 3)
placed = jax.device_put(x, sharding)
# Same global shape and values, placed on the one device in this mesh.

The code explicitly selects one device. Running it on a machine with more devices still uses only that one. To split rows across several devices, construct a larger mesh and use a compatible row count.

Problem

Implement shard_first_axis(x) using the placement pattern above:

  1. Build a one-device mesh from jax.devices()[:1], named 'data'.
  2. Create NamedSharding(mesh, PartitionSpec('data', None)).
  3. Return jax.device_put(x, sharding).

x is a two-dimensional float32 JAX array. Return an array with the same shape, dtype and values and the requested named sharding. This exercise can run on one device and does not require a GPU.

The numerical fixtures verify preserved values, not the returned sharding metadata. They cannot distinguish a correct placement from simply returning x. Inspect the result's .sharding when practicing the API.

Common pitfalls

  • A specification may omit trailing unsplit dimensions. PartitionSpec('data') works for a matrix like PartitionSpec('data', None); two entries are not mandatory. An overly long spec is a different issue.
  • Use Python None, not the string 'None'. Every named mesh axis in the spec must exist in the mesh.
  • Values are preserved on both one-device and multi-device placement. That does not make placement purely metadata: it may move or copy data, and it can change which device owns the array.
  • A one-device mesh teaches the API but cannot demonstrate splitting work across multiple devices. Equal-sized partitions on a larger mesh require dimensions divisible by the assigned mesh-axis sizes.

The diagram below illustrates a separate two-axis layout: rows are split across data and columns across model. The function above uses row-only partitioning on its one-device mesh.

Loading visualization…