Mesh Creation

Why this matters

In JAX’s distributed-computing model, a Mesh is the device topology: a multi-dimensional grid of physical devices (CPUs, GPUs, or TPUs) where each grid axis is given a name. Every sharding strategy you write later — with PartitionSpec and NamedSharding — refers to these axis names to say which dimension of an array maps to which group of devices.

import jax, numpy as np

devices = np.array(jax.devices()[:4]).reshape(2, 2)
mesh = jax.sharding.Mesh(devices, ('batch', 'model'))
# A 2×2 mesh: axis 'batch' has 2 slices, axis 'model' has 2 slices.

The Mesh object owns the mapping from axis names to device slices. len(mesh.axis_names) is the number of named axes (i.e. the number of dimensions in the device grid).

Worked mini-example

import jax
import jax.numpy as jnp
import numpy as np

devices = np.array(jax.devices()[:2])          # 1-D array of 2 devices
mesh = jax.sharding.Mesh(devices, ('data',))   # 1-D mesh, axis named 'data'
print(mesh.axis_names)                          # ('data',)
print(len(mesh.axis_names))                     # 1

Common pitfalls

• Devices must be a NumPy array, not a Python list. Passing a plain list raises a TypeError. Use np.array(jax.devices()[:n]).
• axis_names is a tuple of strings, one per mesh dimension, and its length must equal the rank of the device grid. A 2-D grid needs two names.
• One device is enough for any rank. jax.devices()[:1] reshaped to (1, 1) is a 2-D grid of one device, and Mesh accepts it. The single-CPU test runner is not a limitation on the number of axes.

Problem

Implement mesh_axis_count(n_axes) that:

1. Casts n_axes to int.
2. Takes the one available device and reshapes it into an n-dimensional grid: np.array(jax.devices()[:1]).reshape((1,) * n).
3. Creates a jax.sharding.Mesh over that grid with n axis names — ('axis0',), ('axis0', 'axis1'), and so on.
4. Returns jnp.float32(len(mesh.axis_names)).

Single-device caveat: the test runner has exactly 1 CPU device. A device grid of shape (1, 1) is still a genuine 2-D mesh — every axis simply has one slice — so the whole API works here, and what changes between test cases is the rank of the grid, not the device count.

• n_axes: scalar (cast to int) — how many named axes the mesh should have.

Returns: scalar float32 — the number of mesh axes.

Example (not from the test set):

• mesh_axis_count(2.0) → 2.0

  Loading visualization…