Chapter 5: Parallelization Across Devices with pmap

# Must be set before import jax! Otherwise you need to restart the kernel
import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"

import io, contextlib, time, warnings
warnings.simplefilter("default")  # Make DeprecationWarning and friends visible so we can check API status

import jax
import jax.numpy as jnp
from jax import lax
import numpy as np
from jax.sharding import Mesh, NamedSharding, PartitionSpec as P, AxisType

print("JAX version:", jax.__version__)
print("Devices:", jax.devices())
assert jax.device_count() == 8, "Did not get 8 devices: restart the kernel and run this cell first"

def same(a, b, **kw):
    """Check whether two arrays match: convert to numpy first, since arrays on different meshes can't be combined directly."""
    return np.allclose(np.asarray(a), np.asarray(b), **kw)

# Course example: check the devices
import jax
import jax.numpy as jnp

print(jax.devices())
print("Number of devices:", len(jax.devices()))

# Course example (old API pmap): leading input dimension = number of devices
num_devices = len(jax.devices())

def simple_computation(x):
    return x * 2

pmapped_computation = jax.pmap(simple_computation)
input_data = jnp.arange(num_devices * 3).reshape((num_devices, 3))

try:
    result = pmapped_computation(input_data)
    print("Result:\n", result)
    print("Type:", type(result))
    print("Shape:", result.shape)
    print("Number of devices it lives on:", len(result.devices()))
except Exception as e:
    print("Error:", e)

# ✅ Modern style: jit + NamedSharding, no manual reshape, written just like single-device code
mesh = jax.make_mesh((8,), ("data",))
x_sharded = jax.device_put(jnp.arange(8 * 3), NamedSharding(mesh, P("data")))   # (24,) dimension 0 split over 'data'
print("Local shape per device:", [s.data.shape for s in x_sharded.addressable_shards])

y_jit = jax.jit(simple_computation)(x_sharded)
print("Result sharding:", y_jit.sharding.spec)

# ✅ Modern style: shard_map, which maps one-to-one onto pmap
y_sm = jax.jit(jax.shard_map(simple_computation, mesh=mesh, in_specs=P("data"), out_specs=P("data")))(x_sharded)

assert same(y_jit, result.ravel()) and same(y_sm, result.ravel())
print("✓ jit+sharding, shard_map and pmap give the same result")

x83 = jnp.arange(24).reshape(8, 3)
shapes = {}
jax.pmap(lambda x: (shapes.__setitem__("pmap", x.shape), x)[1])(x83)
jax.shard_map(lambda x: (shapes.__setitem__("shard_map", x.shape), x)[1],
              mesh=mesh, in_specs=P("data"), out_specs=P("data"))(jax.device_put(x83, NamedSharding(mesh, P("data"))))
print("Local shape seen inside the device:", shapes)
assert shapes == {"pmap": (3,), "shard_map": (1, 3)}

with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    out = jax.pmap(lambda x: x * 2)(jnp.arange(8.0))
print("Warnings raised by calling pmap:", len(caught))
print("\n--- start of jax.pmap.__doc__ ---")
print(jax.pmap.__doc__[:640])
print("--- return value ---")
print("Is a jax.Array:", isinstance(out, jax.Array))
print("sharding:", out.sharding)
assert len(caught) == 0 and isinstance(out, jax.Array)
assert "Old way" in jax.pmap.__doc__ and "_internal_pmap_axis_name" in str(out.sharding)

# ❌ Old style (outdated): jax.experimental.shard_map, common in older material, is deprecated since v0.8.0
with warnings.catch_warnings(record=True) as caught:
    warnings.simplefilter("always")
    from jax.experimental.shard_map import shard_map as old_shard_map
print("Warnings on import:", [str(w.message) for w in caught])
assert any(issubclass(w.category, DeprecationWarning) for w in caught)

# ✅ Modern style: use jax.shard_map directly
assert hasattr(jax, "shard_map")
print(jax.shard_map)

# Course example (old API pmap): in_axes=(0, None) means data is mapped along axis 0 and scalar is replicated to every device
import numpy as np

def scale(x, k):
    # Scale x by the scalar factor k
    return x * k

data = jnp.arange(12, dtype=jnp.float32).reshape((4, 3))   # 4 entries, feature size 3
scalar = jnp.float32(10.0)

pmapped_scale = jax.pmap(scale, in_axes=(0, None), out_axes=0)
result = pmapped_scale(data, scalar)

print("Number of devices:", jax.local_device_count())
print("Original data shape:", data.shape)
print("pmap result shape:", result.shape)
print("pmap result:\n", result)
print("Devices that actually took part:", len(result.devices()))
assert len(result.devices()) == 4   # Only 4 devices used: one row (3,) per device

# ❌ Wrong: pmap over a (4, 3) array on 2 devices (devices= simulates "only 2 devices")
two = jax.devices()[:2]
try:
    jax.pmap(scale, in_axes=(0, None), devices=two)(data, scalar)
    raise AssertionError("Should not have succeeded")
except ValueError as e:
    print("Actual error:", str(e)[:120])

# Likewise: a mapped axis of length 16 > 8 devices also errors (the 0.12 error message isn't very clear)
try:
    jax.pmap(lambda x: x * 2)(jnp.arange(16.0))
    raise AssertionError("Should not have succeeded")
except ValueError as e:
    print("16 > 8 error:", str(e)[:120])

# ✅ Fix: reshape to (N, B/N, ...) first, so each device gets a (B/N, ...) per-device batch
def scale_shape(x, k):
    print("  shape of x inside the device:", x.shape)   # printed once, during tracing
    return x * k

per_device = data.reshape(2, 2, 3)                                   # (4, 3) -> (2 devices, 2 rows, 3)
r_fixed = jax.pmap(scale_shape, in_axes=(0, None), devices=two)(per_device, scalar)
assert r_fixed.shape == (2, 2, 3)
assert same(r_fixed.reshape(4, 3), data * scalar)
print("✓ Runs correctly on 2 devices after the reshape")

# ✅ Modern style: jit + NamedSharding. This is the real "cut into B/N blocks automatically" behavior, no reshape needed
mesh2 = jax.make_mesh((2,), ("data",), devices=two)
data_sh = jax.device_put(data, NamedSharding(mesh2, P("data", None)))
print("Local shape per device:", data_sh.addressable_shards[0].data.shape)   # (2, 3): each device gets 2 rows

r_jit = jax.jit(scale)(data_sh, scalar)
assert data_sh.addressable_shards[0].data.shape == (2, 3)
assert same(r_jit, result) and same(r_jit, r_fixed.reshape(4, 3))
print("✓ jit+sharding matches pmap")

x8 = jnp.arange(16.0).reshape(8, 2)

# Old API pmap: pmean first, then out_axes=None (correct usage)
ok = jax.pmap(lambda x: lax.pmean(x.sum(), "d"), axis_name="d", out_axes=None)(x8)

# ❌ Forgot to aggregate + out_axes=None: pmap doesn't complain and quietly returns one device's local value 0+1=1
bad = jax.pmap(lambda x: x.sum(), out_axes=None)(x8)
print("Correct (after pmean):", ok, "   forgot to aggregate:", bad)
assert float(ok) == 15.0                 # Global: the mean of each device's local sum
assert float(bad) == float(x8[0].sum())  # Silently equals device 0's local value, with no error at all

# ✅ Modern style: shard_map's out_specs=P() statically checks "is this really replicated"; without aggregation it errors out
xs16 = jax.device_put(jnp.arange(16.0), NamedSharding(mesh, P("data")))
ok_sm = jax.shard_map(lambda x: lax.pmean(x.sum(), "data"), mesh=mesh, in_specs=P("data"), out_specs=P())(xs16)
assert same(ok_sm, ok)
try:
    jax.shard_map(lambda x: x.sum(), mesh=mesh, in_specs=P("data"), out_specs=P())(xs16)
    raise AssertionError("Should have raised an error")
except ValueError as e:
    print("shard_map error (a good thing):", str(e).splitlines()[0][:110], "...")

# Old API pmap: params replicated as a whole, x mapped along axis 1
params = {"w": jnp.ones((3,)), "b": jnp.float32(0.5)}
xb = jnp.arange(3 * 8, dtype=jnp.float32).reshape(3, 8)   # device dimension on axis 1

def affine(p, x):
    return x * p["w"] + p["b"]

out_p = jax.pmap(affine, in_axes=(None, 1))(params, xb)     # device i gets xb[:, i], shape (3,)
print("pmap output shape:", out_p.shape)                         # (8, 3): device dimension stacked on axis 0

# You can also specify each entry of the dict separately: w replicated, b mapped
params2 = {"w": jnp.ones((3,)), "b": jnp.arange(8.0)}
out_p2 = jax.pmap(affine, in_axes=({"w": None, "b": 0}, 1))(params2, xb)

# ✅ Modern style: jit + NamedSharding. Split dimension 1 of xb over the 'data' axis and replicate the parameters
xb_sh = jax.device_put(xb, NamedSharding(mesh, P(None, "data")))
out_m = jax.jit(lambda p, x: (x * p["w"][:, None] + p["b"]).T)(params, xb_sh)
assert same(out_m, out_p)

# b mapped inside the dict: b is split over 'data' as well
b_sh = jax.device_put(params2["b"], NamedSharding(mesh, P("data")))
out_m2 = jax.jit(lambda w, b, x: x.T * w + b[:, None])(params2["w"], b_sh, xb_sh)
assert same(out_m2, out_p2)
print("✓ Modern versions of both pytree in_axes examples match pmap")

# Course example: list the devices (output depends on the hardware)
print(jax.devices())
# Single CPU: [CpuDevice(id=0)]   Multi-GPU: [CudaDevice(id=0), ...] (shown as GpuDevice in older JAX)   TPU: [TpuDevice(id=0), ...]

# Course example (old API pmap): name the parallel axis with axis_name
def simple_computation(x):
    return x * 2

data = jnp.arange(4 * 5).reshape((4, 5))   # one row per device (4 rows in total)

# pmap implicitly creates a 1-D mesh of size 4; axis 0 of data maps to the device axis named 'data_parallel_axis'
parallel_computation = jax.pmap(simple_computation, axis_name="data_parallel_axis", in_axes=0)
result = parallel_computation(data)
print(result.shape)   # (4, 5)
print(result)

# ✅ Modern style: shard_map, with the axis name written directly on the mesh
mesh4 = jax.make_mesh((4,), ("data_parallel_axis",), devices=jax.devices()[:4])
result_sm = jax.shard_map(simple_computation, mesh=mesh4,
                          in_specs=P("data_parallel_axis"), out_specs=P("data_parallel_axis"))(
    jax.device_put(data, NamedSharding(mesh4, P("data_parallel_axis"))))
assert same(result_sm, result)
print("✓ shard_map matches pmap; mesh:", mesh4)

# ❌ Old style (outdated): jax.experimental.maps was removed in 0.12
try:
    from jax.experimental.maps import Mesh as OldMesh
    raise AssertionError("Import should not have succeeded")
except ModuleNotFoundError as e:
    print("Actual result:", e)

# ✅ Modern style: jax.make_mesh (recommended, arranges devices to match the hardware topology) or jax.sharding.Mesh (manual arrangement)
mesh2d = jax.make_mesh((4, 2), ("data", "model"))
mesh2d_manual = Mesh(np.array(jax.devices()).reshape(4, 2), ("data", "model"))
print(mesh2d)
print("Device layout:\n", mesh2d.devices)
assert mesh2d.shape == mesh2d_manual.shape == {"data": 4, "model": 2}

A = jnp.arange(8 * 4, dtype=jnp.float32).reshape(8, 4)
for spec, local in [(P("data", "model"), (2, 2)), (P("data", None), (2, 4)), (P(None, "model"), (8, 2)),
                    (P(), (8, 4)), (P(("data", "model"), None), (1, 4))]:
    a = jax.device_put(A, NamedSharding(mesh2d, spec))
    assert a.addressable_shards[0].data.shape == local
    print(f"{str(spec):28s} local shape per device {local}   type {jax.typeof(a)}")

mesh_e = jax.make_mesh((8,), ("data",))                              # Explicit by default
mesh_a = jax.make_mesh((8,), ("data",), axis_types=(AxisType.Auto,))  # Auto, set explicitly
assert mesh_e.axis_types == (AxisType.Explicit,)

f = lambda x: lax.psum(x.sum(), "data")
raw = jnp.arange(16.0)        # no device_put, so it's unsharded
try:
    jax.shard_map(f, mesh=mesh_e, in_specs=P("data"), out_specs=P())(raw)
    raise AssertionError("Should have raised an error")
except ValueError as e:
    print("Explicit mesh + unsharded input error:", str(e)[:100], "...")

r1 = jax.shard_map(f, mesh=mesh_e, in_specs=P("data"), out_specs=P())(jax.device_put(raw, NamedSharding(mesh_e, P("data"))))
r2 = jax.shard_map(f, mesh=mesh_a, in_specs=P("data"), out_specs=P())(raw)
assert float(r1) == float(r2) == 120.0
print("✓ Fix 1 (device_put first) and fix 2 (Auto mesh) both give 120")

x2d = jax.device_put(A, NamedSharding(mesh2d, P("data", "model")))
try:
    jax.debug.visualize_array_sharding(x2d)
except Exception as e:
    print("visualize_array_sharding unavailable:", e)

def show_sharding(x):
    """A simple version without rich: list the slice ranges each device holds."""
    print(f"array shape {x.shape}, spec = {x.sharding.spec}")
    for dev, idx in sorted(x.sharding.devices_indices_map(x.shape).items(), key=lambda kv: kv[0].id):
        rng = ", ".join(f"{s.start or 0}:{s.stop if s.stop is not None else n}" for s, n in zip(idx, x.shape))
        print(f"  {dev}: [{rng}]")

show_sharding(x2d)

# ❌ Old style (outdated): nested pmap errors in 0.12
xs = jnp.ones((2, 4))
try:
    jax.pmap(jax.pmap(lambda x: lax.psum(x, "j"), axis_name="j"), axis_name="i")(xs)
    raise AssertionError("Should not have succeeded")
except ValueError as e:
    print("Actual error:", str(e)[:130])

# ✅ Modern style: a (2, 4) mesh + shard_map; collectives can run along just one axis or along several at once
mesh_ij = jax.make_mesh((2, 4), ("i", "j"))
xs_sh = jax.device_put(xs, NamedSharding(mesh_ij, P("i", "j")))
psum_j = jax.shard_map(lambda x: lax.psum(x, "j"), mesh=mesh_ij, in_specs=P("i", "j"), out_specs=P("i", None))
psum_ij = jax.shard_map(lambda x: lax.psum(x, ("i", "j")), mesh=mesh_ij, in_specs=P("i", "j"), out_specs=P())
assert same(psum_j(xs_sh), [[4.0], [4.0]])   # each row adds up four 1s
assert same(psum_ij(xs_sh), [[8.0]])         # adds up all eight 1s
print("✓ sum along j:", psum_j(xs_sh).ravel(), "  sum along (i, j):", psum_ij(xs_sh).ravel())

# Course example (old API pmap): lax.psum
num_devices = 4   # the example is written for 4 devices (this machine has 8; pmap uses only the first 4)

def calculate_local_sum(x):
    local_value = jnp.sum(x)                                  # each device first sums its local data
    total_sum = lax.psum(local_value, axis_name="devices")    # sum across all devices
    return total_sum

data = jnp.arange(num_devices * 3).reshape((num_devices, 3))
# device 0: [0 1 2] -> 3; device 1: [3 4 5] -> 12; device 2: [6 7 8] -> 21; device 3: [9 10 11] -> 30

pmapped_calculate_sum = jax.pmap(calculate_local_sum, axis_name="devices")   # axis_name must match the one in psum
result_psum = pmapped_calculate_sum(data)
print(result_psum)
assert same(result_psum, [66] * 4)   # 3 + 12 + 21 + 30 = 66

# ✅ Modern style: shard_map + psum (and jit automatic parallelism)
mesh4 = jax.make_mesh((4,), ("devices",), devices=jax.devices()[:4])
data_sh = jax.device_put(data, NamedSharding(mesh4, P("devices")))

def local_sum_then_psum(x):          # x: the (1, 3) block on each device
    local = jnp.sum(x)
    print("  type of local:", jax.typeof(local), "  after psum:", jax.typeof(lax.psum(local, "devices")))
    return lax.psum(local, "devices")

psum_one = jax.jit(jax.shard_map(local_sum_then_psum, mesh=mesh4, in_specs=P("devices"), out_specs=P()))(data_sh)
# To get one copy per device like pmap: keep one dimension in the output and use out_specs=P('devices')
psum_each = jax.shard_map(lambda x: lax.psum(jnp.sum(x), "devices")[None],
                          mesh=mesh4, in_specs=P("devices"), out_specs=P("devices"))(data_sh)
psum_jit = jax.jit(jnp.sum)(data_sh)   # jit automatic parallelism: just write the global sum

assert int(psum_one) == int(psum_jit) == 66 and same(psum_each, result_psum)
print("✓ shard_map(out_specs=P()) =", psum_one, "; one copy per device =", psum_each, "; jit =", psum_jit)

# Course example (old API pmap): lax.pmean
def calculate_local_mean(x):
    local_value = jnp.mean(x)
    global_mean = lax.pmean(local_value, axis_name="devices")
    return global_mean

data = jnp.arange(num_devices * 3, dtype=jnp.float32).reshape((num_devices, 3))
# local means: 1.0, 4.0, 7.0, 10.0
pmapped_calculate_mean = jax.pmap(calculate_local_mean, axis_name="devices")
result_pmean = pmapped_calculate_mean(data)
print(result_pmean)
assert same(result_pmean, [5.5] * 4)   # (1 + 4 + 7 + 10) / 4 = 5.5

# ✅ Modern style: shard_map + pmean, and jit automatic parallelism
data_sh = jax.device_put(data, NamedSharding(mesh4, P("devices")))
pmean_sm = jax.jit(jax.shard_map(lambda x: lax.pmean(jnp.mean(x), "devices"),
                                 mesh=mesh4, in_specs=P("devices"), out_specs=P()))(data_sh)
pmean_jit = jax.jit(jnp.mean)(data_sh)
assert float(pmean_sm) == float(pmean_jit) == float(result_pmean[0]) == 5.5
print("✓ shard_map pmean =", pmean_sm, "; jit mean =", pmean_jit)

# Old API pmap: pmax / pmin / all_gather / axis_index / ppermute
n = 8
xv = jnp.array([3.0, 7.0, 1.0, 5.0, 2.0, 8.0, 6.0, 4.0])

def collectives(v):
    return {
        "pmax": lax.pmax(v, "i"),
        "pmin": lax.pmin(v, "i"),
        "all_gather": lax.all_gather(v, "i"),        # every device gets the full length-8 vector
        "axis_index": lax.axis_index("i"),           # current device index 0..7
        "axis_size": lax.psum(1, "i"),               # psum(1) = number of devices on the axis
        "ppermute": lax.ppermute(v, "i", perm=[(j, (j + 1) % n) for j in range(n)]),   # device j sends to j+1
    }

out_pmap = jax.pmap(collectives, axis_name="i")(xv)
for k, v in out_pmap.items():
    print(f"{k:10s}", v if v.ndim == 1 else f"shape {v.shape}, on device 0: {v[0]}")

# ✅ Modern style: shard_map. Each device's local block is (1,); out_specs=P('i') stitches the output back to (8,) for comparison with pmap
mesh_i = jax.make_mesh((8,), ("i",))
xv_sh = jax.device_put(xv, NamedSharding(mesh_i, P("i")))

@jax.jit
@jax.shard_map(mesh=mesh_i, in_specs=P("i"), out_specs=P("i"))
def collectives_sm(v):                       # v: (1,)
    return {
        "pmax": lax.pmax(v, "i"),
        "pmin": lax.pmin(v, "i"),
        "all_gather": lax.all_gather(v, "i", tiled=True)[None, :],   # (1, 8), becomes (8, 8) once stitched back
        "axis_index": lax.axis_index("i")[None],
        "axis_size": jnp.full((1,), lax.axis_size("i")),
        "ppermute": lax.ppermute(v, "i", perm=[(j, (j + 1) % n) for j in range(n)]),
    }

out_sm = collectives_sm(xv_sh)
for k in out_pmap:
    assert same(out_sm[k], out_pmap[k]), k
print("✓ The shard_map versions of all 6 collectives match pmap exactly")

# The jit automatic-parallelism equivalents: write global max/min directly; all_gather just means "make the array replicated"
assert float(jax.jit(jnp.max)(xv_sh)) == 8.0 and float(jax.jit(jnp.min)(xv_sh)) == 1.0
gathered = jax.reshard(xv_sh, NamedSharding(mesh_i, P()))   # passing a bare P() requires with jax.set_mesh(mesh_i): first
assert gathered.sharding.is_fully_replicated and same(gathered, xv)
print("✓ jit: jnp.max/jnp.min work across devices automatically; resharding to a replicated sharding is the same as all_gather")

# ❌ Wrong: jax.device_put(data, devices), trying to spread data over a list of devices
num_devices = 2
devices = jax.local_devices()[:num_devices]
data = np.arange(8.0).reshape(num_devices, -1)   # (2, 4)
try:
    sharded_data = jax.device_put(data, devices)
    raise AssertionError("Should not have succeeded")
except ValueError as e:
    print("Actual error:", str(e).splitlines()[0])
# Why: device_put's second argument is either a single device / Sharding or a pytree matching the data; a list of devices does not "shard" anything

# ✅ Fix: hand the numpy array straight to pmap (it does the distribution itself), and warm up before timing
def complex_computation(x):
    y = jnp.sin(x) * jnp.cos(x)
    z = jnp.tanh(y) + jnp.sqrt(jnp.abs(x))
    return z * 2.0

pmap_complex_computation = jax.pmap(complex_computation, devices=devices)
jit_then_pmap = jax.pmap(jax.jit(complex_computation), devices=devices)

def bench(fn, x, reps=50):
    fn(x).block_until_ready()                  # warm-up: exclude the first compilation time
    t0 = time.perf_counter()
    for _ in range(reps):
        fn(x).block_until_ready()
    return (time.perf_counter() - t0) / reps

result_pj = pmap_complex_computation(data)
print(f"pmap        : {bench(pmap_complex_computation, data) * 1e6:.1f} µs/call")
print(f"pmap(jit(f)): {bench(jit_then_pmap, data) * 1e6:.1f} µs/call (the outer pmap already compiles, so an explicit jit brings no gain; microsecond differences come from dispatch overhead and measurement noise)")
assert same(result_pj, jit_then_pmap(data))

# ✅ Modern style: jit + NamedSharding
mesh_dev2 = jax.make_mesh((2,), ("devices",), devices=devices)
data_sh = jax.device_put(data, NamedSharding(mesh_dev2, P("devices")))   # the correct "sharded device_put"
result_jit = jax.jit(complex_computation)(data_sh)
assert same(result_jit, result_pj, atol=1e-6)
print("✓ jit+sharding matches pmap; output sharding:", result_jit.sharding.spec)

# Course example (old API pmap): pmap + value_and_grad + pmean, one data-parallel gradient computation
def predict(params, x):
    w, b = params
    return w * x + b                     # linear model y = w x + b

def loss_fn(params, x_batch, y_batch):
    predictions = predict(params, x_batch)
    error = predictions - y_batch
    return jnp.mean(error ** 2)          # mean squared error

value_and_grad_fn = jax.value_and_grad(loss_fn)

def parallel_update_step(params, x_shards, y_shards):
    loss, grads = value_and_grad_fn(params, x_shards, y_shards)   # each device computes local gradients
    avg_grads = lax.pmean(grads, axis_name="devices")             # average gradients across devices
    avg_loss = lax.pmean(loss, axis_name="devices")
    return avg_loss, avg_grads

# params replicated (None), x / y mapped along axis 0; both outputs are already aggregated, so out_axes=None
pmap_update_step = jax.pmap(parallel_update_step, axis_name="devices",
                            in_axes=(None, 0, 0), out_axes=(None, None), devices=devices)

params = (jnp.array(2.0), jnp.array(1.0))                 # (w, b)
x_data = np.arange(8.0).reshape(num_devices, -1)          # (2, 4)
# Generate the noise with jax.random.key (not an unseeded np.random.randn) so results are reproducible
noise = jax.random.normal(jax.random.key(0), x_data.shape) * 0.5
y_data = 3 * x_data + 2 + np.asarray(noise)               # y = 3x + 2 + noise

avg_loss, avg_grads = pmap_update_step(params, x_data, y_data)
print(f"Average loss across devices: {avg_loss:.4f}")
print(f"Average gradient (dw, db): ({avg_grads[0]:.4f}, {avg_grads[1]:.4f})")

# Check: matches the single-device full-batch computation
full_loss, full_grads = value_and_grad_fn(params, x_data.ravel(), y_data.ravel())
assert same(avg_loss, full_loss) and same(avg_grads[0], full_grads[0]) and same(avg_grads[1], full_grads[1])

# ✅ Modern style 1: jit + NamedSharding. Call value_and_grad directly on sharded data, with no collectives at all
xs_sh = jax.device_put(x_data.ravel(), NamedSharding(mesh_dev2, P("devices")))
ys_sh = jax.device_put(y_data.ravel(), NamedSharding(mesh_dev2, P("devices")))
l1, g1 = jax.jit(value_and_grad_fn)(params, xs_sh, ys_sh)

# ✅ Modern style 2: shard_map. Differentiate the "global loss pmean(local loss)" (don't copy pmean(grad); see the next special topic)
@jax.jit
@jax.shard_map(mesh=mesh_dev2, in_specs=(P(), P("devices"), P("devices")), out_specs=(P(), P()))
def update_step_sm(params, x, y):
    return jax.value_and_grad(lambda p: lax.pmean(loss_fn(p, x, y), "devices"))(params)

l2, g2 = update_step_sm(params, xs_sh, ys_sh)
for l, g in [(l1, g1), (l2, g2)]:
    assert same(l, avg_loss) and same(g[0], avg_grads[0]) and same(g[1], avg_grads[1])
print("✓ jit+sharding, shard_map and pmap give exactly the same loss and gradients")

# Course example (old API pmap): using vmap inside pmap
def process_item(item):
    return jnp.tanh(item) * 2.0

def process_batch_of_items(batch):      # runs on each device
    vectorized_processor = jax.vmap(process_item)
    return vectorized_processor(batch)

pmap_process_batches = jax.pmap(process_batch_of_items, in_axes=0)

data = np.arange(2 * 3 * 4.0).reshape(num_devices, 3, 4)   # (2, 3, 4)
result_pv = pmap_process_batches(data)
print("Input shape per device:", data.shape[1:])   # (3, 4)
print("Output shape:", result_pv.shape)            # (2, 3, 4)
# pmap distributes along the outermost dimension of size 2; each device gets (3, 4), and the inner vmap vectorizes over the axis of size 3

# ✅ Modern style: jit(vmap(...)) + sharded input. The outer vmap replaces pmap's mapping, and the sharding decides which device holds what
data_sh = jax.device_put(data, NamedSharding(mesh_dev2, P("devices")))
result_mv = jax.jit(jax.vmap(process_batch_of_items))(data_sh)
assert same(result_mv, result_pv)
print("✓ Same result; local shape per device:", result_mv.addressable_shards[0].data.shape)

N_DEV = jax.device_count()
mesh8 = jax.make_mesh((N_DEV,), ("data",))

def loss_w(w, x):
    return jnp.mean((x * w - 1.0) ** 2)

xs = jnp.arange(16.0) / 16
xs_sh = jax.device_put(xs, NamedSharding(mesh8, P("data")))
w0 = jnp.array(0.5)

def smap(f):
    return jax.jit(jax.shard_map(f, mesh=mesh8, in_specs=(P(), P("data")), out_specs=P()))

ref = jax.grad(loss_w)(w0, xs)                                  # single-device global gradient (the correct value)
g_pmap = jax.pmap(lambda w, x: lax.pmean(jax.grad(loss_w)(w, x), "d"),
                  axis_name="d", in_axes=(None, 0))(w0, xs.reshape(N_DEV, -1))[0]
# ❌ Copying the pmap pattern: pmean after grad
g_naive = smap(lambda w, x: lax.pmean(jax.grad(loss_w)(w, x), "data"))(w0, xs_sh)

print(f"single-device correct value: {float(ref):.6f}")
print(f"pmap                       : {float(g_pmap):.6f}")
print(f"shard_map, copied as is    : {float(g_naive):.6f}   ratio = {float(g_naive / ref):.4f}")
assert same(g_pmap, ref)
assert np.isclose(float(g_naive / ref), N_DEV)    # exactly the number of devices

import re
jaxpr_text = str(jax.make_jaxpr(
    jax.shard_map(lambda w, x: jax.grad(loss_w)(w, x), mesh=mesh8, in_specs=(P(), P("data")), out_specs=P())
)(w0, xs_sh))
jaxpr_text = re.sub(r"\x1b\[[0-9;]*m", "", jaxpr_text)   # strip terminal color codes
lines = jaxpr_text.splitlines()
for ln in lines[:3] + ["    ..."] + [l for l in lines if "psum" in l]:
    print(ln)
assert "pvary" in jaxpr_text and "psum_invariant" in jaxpr_text

# ✅ Option 1 (recommended): differentiate the "global loss" pmean(local loss). Clearest meaning, and it doesn't depend on the details of the automatic psum
g1 = smap(lambda w, x: jax.grad(lambda w: lax.pmean(loss_w(w, x), "data"))(w))(w0, xs_sh)

# ✅ Option 2: grad already summed automatically, so divide by the number of devices
g2 = smap(lambda w, x: jax.grad(loss_w)(w, x) / lax.axis_size("data"))(w0, xs_sh)

# ✅ Option 3: first use pcast to make the parameter explicitly varying, restoring pmap's "local gradient" meaning, then pmean
g3 = smap(lambda w, x: lax.pmean(jax.grad(loss_w)(lax.pcast(w, "data", to="varying"), x), "data"))(w0, xs_sh)

# Also: jit + sharding doesn't have this problem at all
g4 = jax.jit(jax.grad(loss_w))(w0, xs_sh)

for name, g in [("Option 1", g1), ("Option 2", g2), ("Option 3", g3), ("jit+sharding", g4)]:
    assert same(g, ref), name
    print(f"✓ {name:12s} {float(g):.6f}")

# ❌ Wrong: the function body uses lax.axis_index, but only `from jax import pmap` was written and lax was never imported
course_src = """
import jax
import jax.numpy as jnp
from jax import pmap

num_devices = 4
sharded_data = jnp.arange(4 * 10).reshape((num_devices, 10))
replicated_data = jnp.array(5.0)

def my_func(x, y):
    return x * y + lax.axis_index('batch')

pmapped_func = pmap(my_func, in_axes=(0, None), out_axes=0, axis_name='batch')
output = pmapped_func(sharded_data, replicated_data)
"""
try:
    exec(course_src, {})          # run this code in a clean namespace
    raise AssertionError("Should not have succeeded")
except NameError as e:
    print("Actual error: NameError -", e)

# ✅ Fix: add from jax import lax
from jax import pmap, lax

num_devices = 4   # in real code, use jax.local_device_count()
sharded_data = jnp.arange(4 * 10).reshape((num_devices, 10))
replicated_data = jnp.array(5.0)

def my_func(x, y):
    # x is sharded (shape (10,) on each device), y is replicated (a scalar)
    return x * y + lax.axis_index("batch")

pmapped_func = pmap(my_func, in_axes=(0, None), out_axes=0, axis_name="batch")

print("Sharded data shape:", sharded_data.shape)      # (4, 10)
assert sharded_data.shape[0] <= jax.local_device_count()
print("Replicated data shape:", replicated_data.shape)   # ()
output_a = pmapped_func(sharded_data, replicated_data)
print("Output shape:", output_a.shape)              # (4, 10)
assert output_a.shape == (4, 10)
assert same(output_a, sharded_data * 5.0 + jnp.arange(4)[:, None])   # row i gets device index i added

# ✅ Modern style: shard_map. The local block is (1, 10); lax.axis_index works the same way
mesh_b = jax.make_mesh((4,), ("batch",), devices=jax.devices()[:4])
out_sm = jax.shard_map(my_func, mesh=mesh_b, in_specs=(P("batch"), P()), out_specs=P("batch"))(
    jax.device_put(sharded_data, NamedSharding(mesh_b, P("batch"))), replicated_data)
assert same(out_sm, output_a)
print("✓ shard_map matches pmap")

# ❌ Wrong: using a Python if on the traced value lax.axis_index(...) == 0
def func_with_debug_print(x):
    intermediate = x * 2
    if lax.axis_index("data_axis") == 0:
        jax.debug.print("device 0 intermediate: {val}", val=intermediate)
    result = lax.psum(intermediate, axis_name="data_axis")
    jax.debug.print("device {id} result after psum: {res}", id=lax.axis_index("data_axis"), res=result)
    return result

data = jnp.arange(2 * 3).reshape((2, 3))     # [[0 1 2], [3 4 5]]
try:
    pmap(func_with_debug_print, axis_name="data_axis")(data)
    raise AssertionError("Should not have succeeded")
except jax.errors.TracerBoolConversionError as e:
    print("Actual error: TracerBoolConversionError -", str(e).splitlines()[0])

# ✅ Fix: use lax.cond to decide whether to print based on the device index
def func_with_debug_print_fixed(x):
    intermediate = x * 2
    # Pass intermediate to the branches explicitly as an operand (don't rely on closure capture), so this also works in shard_map
    lax.cond(lax.axis_index("data_axis") == 0,
             lambda v: jax.debug.print("device 0 intermediate: {val}", val=v),
             lambda v: None,
             intermediate)
    result = lax.psum(intermediate, axis_name="data_axis")
    jax.debug.print("device {id} result after psum: {res}", id=lax.axis_index("data_axis"), res=result)
    return result

output_b = pmap(func_with_debug_print_fixed, axis_name="data_axis")(data)
jax.effects_barrier()   # wait for all asynchronous jax.debug.print output to finish before printing the return value
print("Return value:\n", output_b)
# device 0 intermediate [0 2 4] + device 1 intermediate [6 8 10] = [6 10 14]
assert same(output_b, [[6, 10, 14], [6, 10, 14]])
assert same(output_b[0], data.sum(0) * 2)       # check the expected value with a single-device computation

# ✅ Modern style: jax.debug.print / lax.cond / lax.axis_index work exactly the same way in shard_map
mesh_d = jax.make_mesh((2,), ("data_axis",), devices=jax.devices()[:2])
# Note the outer jax.jit: in 0.12, a shard_map run eagerly op by op without jit raises NotImplementedError when it hits lax.cond
out_dbg = jax.jit(jax.shard_map(lambda x: func_with_debug_print_fixed(x[0])[None],   # local block (1, 3) -> take out (3,)
                                mesh=mesh_d, in_specs=P("data_axis"), out_specs=P("data_axis")))(
    jax.device_put(data, NamedSharding(mesh_d, P("data_axis"))))
assert same(out_dbg, output_b)
print("✓ The shard_map version gives the same result:", out_dbg.tolist())

def show_err(title, fn):
    try:
        fn()
        print(f"[{title}] no error")
        return None
    except Exception as e:
        print(f"[{title}] {type(e).__name__}: {str(e).splitlines()[0][:130]}")
        return e

x8 = jnp.arange(8.0)
errs = [
    show_err("axis name mismatch          ", lambda: jax.pmap(lambda x: lax.psum(x, "dev"), axis_name="devices")(x8)),
    show_err("psum outside a parallel fn  ", lambda: lax.psum(x8, "i")),
    show_err("mapped axis 16 > 8 devices  ", lambda: jax.pmap(lambda x: x)(jnp.arange(16.0))),
    show_err("in_axes=0 on a scalar       ", lambda: jax.pmap(lambda x, y: x * y, in_axes=(0, 0))(x8, 2.0)),
    show_err("Python if on a traced value ", lambda: jax.pmap(lambda x: x if x > 0 else -x)(x8)),
]
assert all(e is not None for e in errs)
assert isinstance(errs[0], NameError) and isinstance(errs[1], NameError)

with jax.disable_jit():
    r = jax.pmap(lambda x: lax.psum(x, "i"), axis_name="i")(jnp.arange(4.0))
assert same(r, [6.0] * 4)
print("\n✓ pmap still runs under disable_jit:", r)

# Check the devices
num_devices = jax.device_count()
print(f"Number of devices available: {num_devices}")
print(f"Available devices: {jax.devices()}")

# Course example 1 (old API pmap): simple data parallelism
def scale_by_two(x):
    print("Compiling and running scale_by_two...")   # runs only during tracing
    return x * 2

data_size_per_device = 4
total_data_size = num_devices * data_size_per_device
global_data = jnp.arange(total_data_size)
sharded_data = global_data.reshape((num_devices, data_size_per_device))
print("Global data shape:", global_data.shape)
print("Sharded data shape:", sharded_data.shape)

parallel_scale = jax.pmap(scale_by_two)
result1 = parallel_scale(sharded_data)
print("\nOutput shape:", result1.shape)
print("Output type:", type(result1))
print("Output content:\n", result1)

# Check: a plain print fires only once across 8 devices; the second call hits the cache and prints 0 times
def scale_by_two_counted(x):
    print("Compiling and running...")
    return x * 2

p = jax.pmap(scale_by_two_counted)
buf1, buf2 = io.StringIO(), io.StringIO()
with contextlib.redirect_stdout(buf1):
    p(sharded_data)
with contextlib.redirect_stdout(buf2):
    p(sharded_data)
print("Prints on the first call:", buf1.getvalue().count("Compiling"), "  second call:", buf2.getvalue().count("Compiling"))
assert buf1.getvalue().count("Compiling") == 1 and buf2.getvalue().count("Compiling") == 0

# ✅ Modern style: jit + NamedSharding, no reshape needed
mesh_dev = jax.make_mesh((num_devices,), ("devices",))
g_sh = jax.device_put(global_data, NamedSharding(mesh_dev, P("devices")))   # (32,), 4 per device
result1_m = jax.jit(lambda x: x * 2)(g_sh)
assert same(result1_m, result1.ravel())
print("✓ Same result; local shape per device:", result1_m.addressable_shards[0].data.shape)

# Course example 2 (old API pmap): in_axes=(0, None)
def scale_and_add(x, y):
    # process_index is the process index (always 0 on one machine), not the device index, so this says "in process" rather than "on device"
    print(f"Compiling and running scale_and_add in process {jax.process_index()}...")
    return x * 2 + y

scalar_y = jnp.float32(100.0)
parallel_scale_add = jax.pmap(scale_and_add, in_axes=(0, None))
result2 = parallel_scale_add(sharded_data, scalar_y)
print("\nInput x shape:", sharded_data.shape, "  Input y:", scalar_y)
print("Output shape:", result2.shape)
print("Output content:\n", result2)    # each element = x * 2 + 100