Skip to content

← JAX Performance & Distributed step 11 of 25

Medium Primitives

Mesh Creation

Jump to the task

Why this matters

A device mesh is a logical grid of JAX devices with a name for each grid axis. Sharding specifications use those names to describe how an array is divided or replicated. The grid is a programming layout; its shape alone does not describe the physical network connecting the devices.

Mesh rank and device count are different. A grid of shape (2, 4) has two axes and eight devices. A grid of shape (1, 1) also has two axes but contains just one device.

Worked mini-example

import jax
import numpy as np
from jax.sharding import Mesh

devices = np.array(jax.devices()[:1]).reshape(1, 1)
mesh = Mesh(devices, ('data', 'model'))
assert mesh.axis_names == ('data', 'model')
assert len(mesh.axis_names) == 2
assert mesh.size == 1

This example runs with one available device. To create a genuine 2-by-2 grid of distinct devices, you need four devices; slicing a shorter device list does not create the missing ones.

Problem

Implement mesh_axis_count(n_axes):

  1. Convert n_axes to a Python integer n.
  2. Reshape the first available device into a grid of shape (1,) * n.
  3. Construct Mesh(devices, tuple(f"axis{i}" for i in range(n))).
  4. Return the length of mesh.axis_names as a float32 scalar.

n_axes is an integer-valued scalar from 1 through 4 and may arrive as a float. Always use just the first device, even if more are available. No array placement or distributed computation is required.

The output tests verify the count. Returning n directly could produce the same output, so the count alone does not prove a mesh was constructed.

Common pitfalls

  • Supply one distinct name per grid dimension. The names label mesh axes, not individual devices or array dimensions.
  • A Python device list is accepted by Mesh. Using a NumPy object array makes a multidimensional grid easy to reshape.
  • Size-one mesh axes are still axes. Count axis_names, not devices.
  • Mesh construction does not itself distribute a tensor. A sharding specification and array placement are separate steps.
  • The Python integer conversion and NumPy reshape require concrete n. Mark n_axes static if wrapping this setup exercise in jax.jit.

The diagram below is a separate eight-device example, not the one-device mesh used by the function. Its array is replicated across the mesh.

Loading visualization…