The mesh and the spec

Three different lines below build the same mesh, and JAX agrees they are the same. jax.make_mesh((4, 2), ("data", "model")) is the short form; a Mesh wrapped around mesh_utils.create_device_mesh is the older one; a Mesh over a plain reshaped array of jax.devices() is what both come down to. A mesh has no identity past the devices it holds and the names it gives their axes.

That equality is worth more than it first looks. A NamedSharding built against a mesh inside a helper compares equal to one built at the call site, so passing meshes around does not quietly leave you with two shardings that disagree about the same layout.

Everything about a mesh reads back. mesh.shape is an ordered mapping from axis name to length, mesh.axis_names is the tuple you passed, and mesh.size is the device count. No array has appeared yet, which is the point: a mesh is a grid of devices, and you can build one before deciding what to put on it.

import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"

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

a = jax.make_mesh((4, 2), ("data", "model"))
b = Mesh(mesh_utils.create_device_mesh((4, 2)), ("data", "model"))
c = Mesh(np.array(jax.devices()).reshape(4, 2), ("data", "model"))

print(a)
print(a.shape, a.axis_names, a.size)
print(a == b, a == c)

# Mesh('data': 4, 'model': 2)
# OrderedDict({'data': 4, 'model': 2}) ('data', 'model') 8
# True True

The entries of a PartitionSpec line up with the axes of the array, in order. Entry i names the mesh axis that splits array axis i, None leaves that axis whole on every device, and a tuple of names splits one array axis across several mesh axes at once.

shard_shape answers the question without allocating anything. Hand it a global shape and it returns the shape one device ends up holding, so (8, 16) under P("data", "model") gives every device a (2, 8) block, and P(("data", "model"), None) splits the first axis eight ways for a (1, 16) block instead.

A spec shorter than the array's rank is legal, and the entries you left off are read as None. P("data") on a two-axis array means P("data", None), which is why the last two rows of the run below print the same shape. A spec longer than the rank is not legal, and the fourth section has the sentence it raises.

The spec is indexed by array axis. The names inside it are mesh axes.

import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"

import jax
from jax.sharding import NamedSharding, PartitionSpec as P

mesh = jax.make_mesh((4, 2), ("data", "model"))

for spec in [P("data", "model"), P("data", None), P(None, "model"), P(),
             P(("data", "model"), None), P("model", "data"), P("data")]:
    print(spec, NamedSharding(mesh, spec).shard_shape((8, 16)))

# PartitionSpec('data', 'model') (2, 8)
# PartitionSpec('data', None) (2, 16)
# PartitionSpec(None, 'model') (8, 8)
# PartitionSpec() (8, 16)
# PartitionSpec(('data', 'model'), None) (1, 16)
# PartitionSpec('model', 'data') (4, 4)
# PartitionSpec('data',) (2, 16)

.addressable_shards turns the arrangement into eight objects you can look at. Each one carries the device it sits on, the index tuple naming which slice of the global array it is, and the local array itself.

The index is an ordinary tuple of Python slices, so a shard's place in the global array is something you read rather than something you remember. Device 0 holds rows 0 to 2 and columns 0 to 8, device 1 holds the same rows and the next eight columns, and device 2 has moved down a row block.

Printing the first element of each shard against a flat arange confirms the tiling: 0, then 8, then 32. Device order runs along the last mesh axis first, so model advances every device and data advances every two.

Chapter 10 draws this same layout as a picture with jax.debug.visualize_array_sharding. The shard list is the version you can assert on in a test, which is the reason to know both.

import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"

import jax
import jax.numpy as jnp
from jax.sharding import NamedSharding, PartitionSpec as P

mesh = jax.make_mesh((4, 2), ("data", "model"))
x = jax.device_put(jnp.arange(128.).reshape(8, 16), NamedSharding(mesh, P("data", "model")))

print(x.shape, len(x.addressable_shards))
for sh in x.addressable_shards[:3]:
    print(sh.device, sh.index, sh.data.shape, sh.data[0, 0].item())

# (8, 16) 8
# TFRT_CPU_0 (slice(0, 2, None), slice(0, 8, None)) (2, 8) 0.0
# TFRT_CPU_1 (slice(0, 2, None), slice(8, 16, None)) (2, 8) 8.0
# TFRT_CPU_2 (slice(2, 4, None), slice(0, 8, None)) (2, 8) 32.0

Each of these messages names the quantity that failed, which makes them worth reading instead of pattern-matching. An axis that does not divide gives you the divisor and the size in one clause: dimension 0 should be divisible by 4, but it is equal to 6.

A name that is not in the mesh gets caught before any data moves, and the message prints the mesh's axis names so the typo is visible next to what you meant. A repeated name is rejected as well, since one mesh axis cannot split two array axes at once; the devices it names would have to be in two places.

The fourth is the rank rule from two sections up, seen from the failing side. A three-entry spec applied to a two-axis array reports that the sharding is only valid for values of rank at least 3. Short specs pad on the right; long ones raise.

import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"

import jax
import jax.numpy as jnp
from jax.sharding import NamedSharding, PartitionSpec as P

mesh = jax.make_mesh((4, 2), ("data", "model"))
cube = jax.make_mesh((2, 2, 2), ("a", "b", "c"))

cases = [(mesh, P("data", "model"), (6, 16)), (mesh, P("batch", None), (8, 16)),
         (mesh, P("data", "data"), (8, 16)), (cube, P("a", "b", "c"), (8, 16))]
for m, spec, shape in cases:
    try:
        jax.device_put(jnp.ones(shape), NamedSharding(m, spec))
    except ValueError as err:
        print(err)
        print()

# One of device_put args was given the sharding of NamedSharding(mesh=Mesh('data': 4,
# 'model': 2), spec=PartitionSpec('data', 'model'), memory_kind=unpinned_host), which
# implies that the global size of its dimension 0 should be divisible by 4, but it is
# equal to 6 (full shape: (6, 16))
#
# Resource axis: batch of PartitionSpec('batch', None) is not found in mesh:
# ('data', 'model').
#
# A single NamedSharding spec specification can map every mesh axis to at most one
# positional dimension, but PartitionSpec('data', 'data') has duplicate entries for
# `data`
#
# One of device_put args is incompatible with its sharding annotation
# NamedSharding(mesh=Mesh('a': 2, 'b': 2, 'c': 2), spec=PartitionSpec('a', 'b', 'c'),
# memory_kind=unpinned_host): Sharding NamedSharding(mesh=Mesh('a': 2, 'b': 2,
# 'c': 2), spec=PartitionSpec('a', 'b', 'c'), memory_kind=unpinned_host) is only
# valid for values of rank at least 3, but was applied to a value of rank 2.

On this version a PartitionSpec inherits from tuple, so len, indexing and equality against a plain tuple all work, and a lot of code leans on that without meaning to.

The changelog records two steps away from it. Release 0.6.1 stopped PartitionSpec inheriting from a tuple, and 0.10.0 stopped it reporting itself equal to one, with the instruction to convert tuples to specs before comparing. Anything that stores a spec as a tuple, or asserts against a tuple in a test, is the code that will break on the upgrade rather than the code that uses specs normally.

from jax.sharding import PartitionSpec as P

spec = P("data", None)
print(isinstance(spec, tuple), len(spec), spec[0], spec[1])
print(spec == ("data", None))

# True 2 data None
# True