# Arrays, meshes, and sharding
Phase 09: Distributed training · about 80 minutes · 4 logical CPU devices
## What you will be able to do
- Translate a PartitionSpec into local shard shapes.
- Inspect shard indices and verify values.
- Explain communication needs of a global reduction.
- Repair uneven batch statistics with padding and masks.
## The problem
We can split examples across devices, split features across devices, or keep a full copy on each one. Let’s make those choices visible on four logical CPU devices. You’ll inspect the pieces and see why some reductions must combine results. No cluster is needed, and these observations do not measure accelerator performance.
## The idea
Sharding places one global array across a named mesh. Global shape, local shard shape and replication describe different views of the same computation. Make ownership explicit before reasoning about communication.
## Compare the global array with each local piece
A four-by-four array on two devices can be split into two-by-four row shards or four-by-two column shards. Both layouts contain the same global values. Their local operations and communication requirements differ.
Replication gives each participant the full value instead of a disjoint portion. Use explicit labels so duplicated ownership cannot be mistaken for partitioning.
Connect each PartitionSpec entry to its array axis and named mesh axis. Inspect actual addressable shards and compare reconstructed values with a global reference. Logical CPU devices validate this placement exercise without reproducing a multi-host network.
$$
\text{Mesh}(D, (\texttt{'data'}, \texttt{'model'})) \implies \operatorname{shard}(X) \in \mathbb{R}^{(B/P_d) \times (H/P_m)}
$$
### Pause and reason
Must changing row sharding to column sharding change the answer?
Compare your reasoning
No. A correct compatible computation preserves the intended global result. Storage and communication can change even when the result agrees.
Reference solution
```python
# Exercise solution: Compare row and column layouts of a (12,8) array, and repair a ten-row...
# Construct and reshape `other` into the target tensor dimensions.
other = np.arange(96,dtype=np.float32).reshape(12,8)
for spec, expected in [(rows,(3,8)),(columns,(12,2))]:
# Place `placed` explicitly onto the target JAX device.
placed = jax.device_put(other,spec)
for shard in placed.addressable_shards:
# Check tensor shape invariant: `shard.data.shape == expected`
assert shard.data.shape == expected
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_array_equal(np.asarray(shard.data),other[shard.index])
# Construct and reshape `original` into the target tensor dimensions.
original = np.arange(40,dtype=np.float32).reshape(10,4)
# Run the boundary check and catch the expected exception:
try:
jax.device_put(original, rows)
except ValueError:
print("Expected indivisible batch")
else:
raise AssertionError("Expected divisibility failure")
# Place `padded` explicitly onto the target JAX device.
padded = jax.device_put(np.pad(original,((0,2),(0,0))),rows)
# Configure multi-device placement / sharding specification (`mask_sharding`).
mask_sharding = NamedSharding(mesh,P("data"))
mask = jax.device_put(np.array([1]*10+[0]*2,dtype=np.float32),mask_sharding)
# Wrap with `jax.jit` (`masked_mean`) so XLA traces and compiles the function.
masked_mean = jax.jit(lambda a,m:(a*m[:,None]).sum(axis=0)/m.sum(),in_shardings=(rows,mask_sharding),out_shardings=replicated)(padded,mask)
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_allclose(np.asarray(masked_mean),original.mean(axis=0),rtol=1e-6)
assert not np.allclose(np.asarray(padded.mean(axis=0)), original.mean(axis=0))
```
Hint
Four devices divide twelve rows or eight columns.
Reference solution and reasoning
```python
# Change both dimensions (Challenge): The row layout owns three complete feature vectors per...
# Construct and reshape `other` into the target tensor dimensions.
other = np.arange(96,dtype=np.float32).reshape(12,8)
for spec, expected in [(rows,(3,8)),(columns,(12,2))]:
# Place `placed` explicitly onto the target JAX device.
placed = jax.device_put(other,spec)
for shard in placed.addressable_shards:
# Check tensor shape invariant: `shard.data.shape == expected`
assert shard.data.shape == expected
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_array_equal(np.asarray(shard.data),other[shard.index])
```
The row layout owns three complete feature vectors per device; the column layout owns two features for every observation.
Hint
The divisor is the sum of the mask, not padded length.
Reference solution and reasoning
```python
# Pad an uneven batch without corrupting its mean (Challenge): Padding repairs the shape but adds artificial rows.
# Construct and reshape `original` into the target tensor dimensions.
original = np.arange(40,dtype=np.float32).reshape(10,4)
# Run the boundary check and catch the expected exception:
try:
jax.device_put(original, rows)
except ValueError:
print("Expected indivisible batch")
else:
raise AssertionError("Expected divisibility failure")
# Place `padded` explicitly onto the target JAX device.
padded = jax.device_put(np.pad(original,((0,2),(0,0))),rows)
# Configure multi-device placement / sharding specification (`mask_sharding`).
mask_sharding = NamedSharding(mesh,P("data"))
mask = jax.device_put(np.array([1]*10+[0]*2,dtype=np.float32),mask_sharding)
# Wrap with `jax.jit` (`masked_mean`) so XLA traces and compiles the function.
masked_mean = jax.jit(lambda a,m:(a*m[:,None]).sum(axis=0)/m.sum(),in_shardings=(rows,mask_sharding),out_shardings=replicated)(padded,mask)
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_allclose(np.asarray(masked_mean),original.mean(axis=0),rtol=1e-6)
assert not np.allclose(np.asarray(padded.mean(axis=0)), original.mean(axis=0))
```
Padding repairs the shape but adds artificial rows. The mask keeps them out of both numerator and denominator. Checking only divisibility would miss this statistical error.
Answer and explanation
Each device initially holds only some rows.
The reduction spans the partitioned row dimension, so partial contributions must be combined.