# Practice with four virtual CPU devices
Phase 00: Setup & first steps · about 60 minutes · 4 logical CPU devices
## What you will be able to do
- Start four virtual CPU devices in a fresh process.
- Explain global shape, shard index and local shape.
- Compare row partitioning with replication.
- Diagnose divisibility and distinguish CPU practice from TPU evidence.
## The problem
You can explore how JAX splits an array across devices using just your laptop. We’ll create four logical CPU devices, give each one a piece of an array, and compare the result with NumPy. Start with the setup and array/device lessons. These logical devices share your CPU: they help us learn placement, but they do not reproduce accelerator speed or memory.
## The idea
Logical CPU devices let us practice how a global array is divided and addressed. They share a physical host. Treat them as a way to make placement visible, then use a real accelerator run to investigate accelerator memory, communication and speed.
## Distinguish ownership from extra computing power
Suppose a global array has four rows and each logical device owns one row. The global shape still describes all four rows; a local shard describes the one row stored for that owner. Asking for the global shape and inspecting a local buffer are different operations.
The ownership heatmap assigns a category to each device. A darker device color does not mean more work or a faster device. Follow the row boundary and the device label together.
Set the logical-device configuration before JAX initializes its backend, then inspect the actual device count in a fresh process. Changing a flag after initialization does not retroactively repartition an already-running process.
$$
X_{\text{global}} \in \mathbb{R}^{D_0 \times D_1} \;\xrightarrow{\text{shard}}\; \{X^{(d)} \in \mathbb{R}^{(D_0 / P) \times D_1}\}_{d=0}^{P-1}
$$
### Pause and reason
A four-device CPU exercise passes. What can you now claim about four TPU chips?
Compare your reasoning
You have checked the array-placement exercise on logical CPU devices. TPU execution, memory behavior and performance still require a separate run on that hardware. The ownership reasoning transfers; the measured performance does not.
Reference solution
```python
# Exercise solution: Predict and verify the per-device shapes for twelve rows, then...
# Construct and reshape `twelve` into the target tensor dimensions.
twelve = np.arange(24,dtype=np.float32).reshape(12,2)
# Place `z` explicitly onto the target JAX device.
z = jax.device_put(twelve, rows)
# Wrap with `jax.jit` (`z2`) so XLA traces and compiles the function.
z2 = jax.jit(lambda a:a*a, in_shardings=rows, out_shardings=rows)(z)
# Check tensor shape invariant: `all(s.data.shape == (3,2) for s in z2.addressable_shards)`
assert all(s.data.shape == (3,2) for s in z2.addressable_shards)
for shard in z2.addressable_shards:
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_array_equal(np.asarray(shard.data), (twelve*twelve)[shard.index])
# Construct and reshape `uneven` into the target tensor dimensions.
uneven = np.arange(10,dtype=np.float32).reshape(5,2)
# Run the boundary check and catch the expected exception:
try:
jax.device_put(uneven, rows)
except ValueError:
print("Expected indivisible row dimension")
else:
raise AssertionError("Expected divisibility failure")
# Place `fixed` explicitly onto the target JAX device.
fixed = jax.device_put(uneven, replicated)
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_array_equal(np.asarray(fixed), uneven)
assert fixed.is_fully_replicated
```
Hint
Use each shard index to select the corresponding NumPy reference slice.
Reference solution and reasoning
```python
# Move from four rows to twelve (Challenge): The twelve rows split into four groups of three.
# Construct and reshape `twelve` into the target tensor dimensions.
twelve = np.arange(24,dtype=np.float32).reshape(12,2)
# Place `z` explicitly onto the target JAX device.
z = jax.device_put(twelve, rows)
# Wrap with `jax.jit` (`z2`) so XLA traces and compiles the function.
z2 = jax.jit(lambda a:a*a, in_shardings=rows, out_shardings=rows)(z)
# Check tensor shape invariant: `all(s.data.shape == (3,2) for s in z2.addressable_shards)`
assert all(s.data.shape == (3,2) for s in z2.addressable_shards)
for shard in z2.addressable_shards:
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_array_equal(np.asarray(shard.data), (twelve*twelve)[shard.index])
```
The twelve rows split into four groups of three. Index-based comparisons establish that the right values occupy each piece, rather than checking only a total.
Hint
Do not drop a row to hide the error.
Reference solution and reasoning
```python
# Repair a batch that does not divide evenly (Challenge): Five rows cannot be evenly split across the four-way data axis.
# Construct and reshape `uneven` into the target tensor dimensions.
uneven = np.arange(10,dtype=np.float32).reshape(5,2)
# Run the boundary check and catch the expected exception:
try:
jax.device_put(uneven, rows)
except ValueError:
print("Expected indivisible row dimension")
else:
raise AssertionError("Expected divisibility failure")
# Place `fixed` explicitly onto the target JAX device.
fixed = jax.device_put(uneven, replicated)
# Convert `` to a host NumPy array for inspection or verification.
np.testing.assert_array_equal(np.asarray(fixed), uneven)
assert fixed.is_fully_replicated
```
Five rows cannot be evenly split across the four-way data axis. Replication preserves all observations but changes the memory/communication plan; padding requires a mask when computing means.
Answer and explanation
Each device holds a $(3,2)$ slice and the global array remains $(12,2)$.
Partitioning divides the mapped dimension; it does not change the global shape or create accelerator hardware.