Chapter 6: Large-Scale Model Training

# Must be set before importing jax: simulate 8 devices on the CPU for the data-parallel / model-parallel / FSDP demos
import os
os.environ["XLA_FLAGS"] = "--xla_force_host_platform_device_count=8"

import sys, time, re, warnings, tempfile
from functools import partial

import numpy as np
import jax
import jax.numpy as jnp
from jax import lax
from jax.sharding import NamedSharding, PartitionSpec as P, AxisType, reshard
import optax
import flax
from flax import nnx

print("Python", sys.version.split()[0])
print("JAX   ", jax.__version__, "| Flax", flax.__version__, "| Optax", optax.__version__)
print("Devices", jax.devices())
assert jax.device_count() == 8

# ---------- Small helpers used throughout this chapter ----------
def tree_bytes(tree):
    """Total bytes of all array leaves in a PyTree."""
    return sum(x.nbytes for x in jax.tree.leaves(tree) if hasattr(x, "nbytes"))

def max_diff(a, b):
    """Largest elementwise absolute difference between two PyTrees of the same structure (compared in float64)."""
    la, lb = jax.tree.leaves(a), jax.tree.leaves(b)
    assert len(la) == len(lb)
    return max(float(np.max(np.abs(np.asarray(x, np.float64) - np.asarray(y, np.float64)))) for x, y in zip(la, lb))

def bench(f, *args, n=10):
    """Warm up first (including compilation), then time n runs and return the average in milliseconds."""
    jax.block_until_ready(f(*args))
    t0 = time.perf_counter()
    for _ in range(n):
        out = f(*args)
    jax.block_until_ready(out)
    return (time.perf_counter() - t0) / n * 1e3

def mem_of(f, *args):
    """Compile jit(f) and return XLA's memory analysis (without running it)."""
    return jax.jit(f).lower(*args).compile().memory_analysis()

def collectives(compiled_text):
    """Count how often each kind of collective appears in the compiled HLO."""
    ops = ["all-reduce", "all-gather", "reduce-scatter", "all-to-all", "collective-permute"]
    return {op: len(re.findall(op + r"(?:-start)?\(", compiled_text)) for op in ops}

MiB = 2 ** 20

GB = 1e9
n_params = 10e9                      # example: 10 billion parameters

mem = {
    "params (fp32)":      n_params * 4,
    "grads (fp32)":       n_params * 4,
    "Adam m+v (fp32)":    n_params * 2 * 4,
}
for k, v in mem.items():
    print(f"{k:18s} {v / GB:6.0f} GB")
static = mem["params (fp32)"] + mem["Adam m+v (fp32)"]
print(f"{'static (params+optimizer)':14s} {static / GB:6.0f} GB")
print(f"{'total (no activations)':15s} {sum(mem.values()) / GB:6.0f} GB")

assert mem["params (fp32)"] / GB == 40 and mem["Adam m+v (fp32)"] / GB == 80
assert static / GB == 120 and sum(mem.values()) / GB == 160

# Measure with Optax: Adam state is exactly 2x the parameter bytes (plus one int32 step count)
toy = {"w": jnp.zeros((1024, 512)), "b": jnp.zeros(512)}
adam_state = optax.adam(1e-3).init(toy)
ratio = tree_bytes(adam_state) / tree_bytes(toy)
print(f"\nOptax Adam state / params byte ratio = {ratio:.6f}")
assert abs(ratio - 2.0) < 1e-4

def bytes_per_param(mixed_precision=True):
    """Bytes needed per parameter for Adam training (excluding activations)."""
    if mixed_precision:
        return dict(bf16_params=2, bf16_grads=2, fp32_master=4, adam_m=4, adam_v=4)
    return dict(fp32_params=4, fp32_grads=4, adam_m=4, adam_v=4)

mp = bytes_per_param(True)
print("Mixed-precision Adam:", mp, "total", sum(mp.values()), "B/param →", n_params * sum(mp.values()) / GB, "GB")
assert sum(mp.values()) == 16
for n_dev in [8, 64, 512]:
    print(f"  FSDP/ZeRO-3 across {n_dev:3d} devices: {n_params * 16 / n_dev / GB:7.2f} GB per device")

# Training compute: C ≈ 6 · Φ · T
tokens = 1e12
C_total = 6 * n_params * tokens
chip_flops, mfu, n_chips = 300e12, 0.4, 256        # assumptions: 300 TFLOP/s per chip, 40% utilization, 256 chips
days = C_total / (chip_flops * mfu * n_chips) / 86400
print(f"\n10B params × 1T tokens ≈ {C_total:.1e} FLOPs, about {days:.1f} days on 256 chips")

# Per-device traffic of a ring all-reduce
S = n_params * 4
for N in [8, 64]:
    print(f"N={N:3d}: each device sends/receives ≈ {2 * (N - 1) / N * S / GB:.1f} GB of gradient data per step")

# ❌ Old style (outdated): flax.linen is the old API (NNX is recommended for new projects); PRNGKey and jax.tree_util.tree_map are old style too
try:
    import flax.linen as nn

    class SimpleMLP(nn.Module):
        features: list[int]          # list of layer sizes, e.g. [128, 64, 10]

        @nn.compact                  # lets you define submodules inline in __call__
        def __call__(self, x):
            for i, feat in enumerate(self.features):
                x = nn.Dense(features=feat, name=f'dense_{i}')(x)
                if i != len(self.features) - 1:   # ReLU on every layer except the last
                    x = nn.relu(x)
            return x

    key = jax.random.PRNGKey(0)
    dummy_input = jnp.ones((1, 28 * 28))        # one flattened MNIST image (batch=1)
    linen_model = SimpleMLP(features=[128, 10])
    linen_params = linen_model.init(key, dummy_input)['params']
    print("Parameter PyTree structure:", jax.tree_util.tree_map(lambda x: x.shape, linen_params))
    linen_out = linen_model.apply({'params': linen_params}, dummy_input)
    print("Output shape:", linen_out.shape)
    print("Measured: Linen still runs in Flax", flax.__version__, "(old API, not deprecated)")
except Exception as e:
    print("Measured error:", type(e).__name__, e)

# ✅ Recommended: Flax NNX
class SimpleMLP(nnx.Module):
    def __init__(self, din: int, features: list[int], *, rngs: nnx.Rngs):
        dims = [din] + list(features)
        # A list of submodules must be wrapped in nnx.List so NNX recognizes it as a graph node
        self.layers = nnx.List([nnx.Linear(a, b, rngs=rngs) for a, b in zip(dims[:-1], dims[1:])])

    def __call__(self, x):
        for i, layer in enumerate(self.layers):
            x = layer(x)
            if i != len(self.layers) - 1:
                x = nnx.relu(x)
        return x

model = SimpleMLP(28 * 28, [128, 10], rngs=nnx.Rngs(0))     # constructing it initializes it
x_demo = jnp.ones((1, 28 * 28))
out = model(x_demo)                                           # call it directly
print("Output shape:", out.shape)

params_state = nnx.state(model, nnx.Param)                    # pull out every nnx.Param
shapes = jax.tree.map(jnp.shape, nnx.to_pure_dict(params_state))
print("Parameter structure:", shapes)

# Parameter shapes match the Linen version above one to one (Linen's dense_i ↔ NNX's layers[i])
assert shapes["layers"][0]["kernel"] == linen_params["dense_0"]["kernel"].shape == (784, 128)
assert shapes["layers"][1]["kernel"] == linen_params["dense_1"]["kernel"].shape == (128, 10)
assert out.shape == linen_out.shape == (1, 10)

# ❌ Old style (outdated): Haiku has been in maintenance mode since 2023-07, and DeepMind recommends Flax for new projects; haiku isn't installed here
try:
    import haiku as hk

    class HkMLP(hk.Module):
        def __init__(self, features, name=None):
            super().__init__(name=name)
            self.features = features
        def __call__(self, x):
            for i, feat in enumerate(self.features):
                x = hk.Linear(output_size=feat, name=f'linear_{i}')(x)
                if i != len(self.features) - 1:
                    x = jax.nn.relu(x)
            return x

    model = hk.transform(lambda x: HkMLP([128, 10])(x))
    params = model.init(jax.random.PRNGKey(0), jnp.ones((1, 784)))
    print(model.apply(params, jax.random.PRNGKey(0), jnp.ones((1, 784))).shape)
except ImportError as e:
    print("Measured:", type(e).__name__, "-", e)

# ✅ Recommended: Haiku's init / apply correspond to nnx.split / nnx.merge in NNX
graphdef, state = nnx.split(model)          # structure + data
print("GraphDef type:", type(graphdef).__name__, "| State leaf count:", len(jax.tree.leaves(state)))

@jax.jit                                     # plain jax.jit: only accepts PyTrees
def pure_apply(state, x):                    # equivalent to Linen's model.apply(params, x)
    return nnx.merge(graphdef, state)(x)

@nnx.jit                                     # NNX's lifted jit: you can pass the module object directly
def nnx_apply(model, x):
    return model(x)

y1, y2, y3 = pure_apply(state, x_demo), nnx_apply(model, x_demo), model(x_demo)
assert max_diff(y1, y3) == 0.0 and max_diff(y2, y3) == 0.0
print("jax.jit + split/merge, nnx.jit and a direct call all give identical results ✓")

# Merging back gives a new object that is independent of the original
model_copy = nnx.merge(graphdef, state)
assert model_copy is not model and max_diff(nnx.state(model_copy), nnx.state(model)) == 0.0

class Counter(nnx.Module):
    def __init__(self):
        self.count = nnx.Variable(jnp.array(0))

c = Counter()
with warnings.catch_warnings(record=True) as w:
    warnings.simplefilter("always")
    _ = c.count.value                              # old style
dep = [str(m.message) for m in w if issubclass(m.category, DeprecationWarning)]
print("Warning when reading .value:", dep[0].splitlines()[0] if dep else None)
assert dep and "deprecated" in dep[0]

c.count[...] = c.count[...] + 1                    # ✅ new style
assert int(c.count[...]) == 1

# wrt is a required keyword argument of nnx.Optimizer
import inspect
sig = inspect.signature(nnx.Optimizer.__init__)
print("nnx.Optimizer signature:", sig)
assert sig.parameters["wrt"].default is inspect.Parameter.empty
try:
    nnx.Optimizer(model, optax.sgd(0.1))
except TypeError as e:
    print("Without wrt →", type(e).__name__, ":", e)

# Course example (written in the recommended style: jax.random.key, jax.tree.*)
params = {
    'encoder': {
        'layer_1': {'w': jnp.ones((128, 256)), 'b': jnp.zeros(256)},
        'layer_norm': {'scale': jnp.ones(256), 'bias': jnp.zeros(256)},
    },
    'decoder': {
        'output': {'w': jnp.ones((256, 10)), 'b': jnp.zeros(10)},
    },
}
optimizer = optax.adam(learning_rate=1e-3)
opt_state = optimizer.init(params)          # structure mirrors the parameters
key = jax.random.key(0)

training_state = {'params': params, 'opt_state': opt_state, 'rng_key': key, 'step': 0}
leaves, treedef = jax.tree.flatten(training_state)
print("Number of leaves:", len(leaves))

# Break it down item by item
n_p   = len(jax.tree.leaves(params))
n_opt = len(jax.tree.leaves(opt_state))
print(f"params {n_p} + Adam state {n_opt} (count 1 + mu {n_p} + nu {n_p}, EmptyState 0) + key 1 + step 1")
print("opt_state structure:", [type(s).__name__ for s in opt_state])
assert n_p == 6 and n_opt == 1 + 2 * n_p
assert len(leaves) == 21

# ❌ Old style (outdated): Linen + flax.training.train_state is the old API; the NNX counterpart is nnx.Optimizer
try:
    from typing import Any
    import flax.linen as nn
    from flax.training import train_state

    class SimpleModel(nn.Module):
        features: int
        @nn.compact
        def __call__(self, x):
            return nn.Dense(features=self.features)(x)

    class TrainState(train_state.TrainState):
        batch_stats: Any = None              # optional: BatchNorm statistics

    key = jax.random.PRNGKey(0)
    lin_model = SimpleModel(features=10)
    lin_params = lin_model.init(key, jnp.ones([1, 128]))['params']
    ts = TrainState.create(apply_fn=lin_model.apply, params=lin_params, tx=optax.adam(1e-3))
    grads = jax.tree.map(jnp.ones_like, ts.params)
    ts2 = ts.apply_gradients(grads=grads)    # immutable update: returns a new object
    print("Measured: it runs; step", int(ts.step), "→", int(ts2.step), "; the old object is unchanged")
except Exception as e:
    print("Measured error:", type(e).__name__, e)

# ✅ Recommended: in NNX, nnx.Optimizer takes over the role of TrainState (step count and optimizer state live inside it, updated in place)
class SimpleModelNNX(nnx.Module):
    def __init__(self, din, features, *, rngs):
        self.dense = nnx.Linear(din, features, rngs=rngs)
    def __call__(self, x):
        return self.dense(x)

nnx_model = SimpleModelNNX(128, 10, rngs=nnx.Rngs(0))
nnx_opt = nnx.Optimizer(nnx_model, optax.adam(1e-3), wrt=nnx.Param)
grads = jax.tree.map(jnp.ones_like, nnx.state(nnx_model, nnx.Param))
w_old = nnx_model.dense.kernel[...]
nnx_opt.update(nnx_model, grads)                 # equivalent to state.apply_gradients(grads=grads)
print("nnx.Optimizer.step:", int(nnx_opt.step[...]), "| parameters updated in place:", bool(jnp.any(nnx_model.dense.kernel[...] != w_old)))
assert int(nnx_opt.step[...]) == 1 and int(ts2.step) == 1
# Same effect as one update step of the Linen TrainState (same gradients, same Adam; the first update is -lr·sign(g) in both)
assert np.allclose(nnx_model.dense.kernel[...] - w_old, ts2.params["Dense_0"]["kernel"] - ts.params["Dense_0"]["kernel"], atol=1e-7)

class StepCounter(nnx.Variable):          # custom variable type
    pass

class BNNet(nnx.Module):
    def __init__(self, din, dh, dout, *, rngs: nnx.Rngs):
        self.l1 = nnx.Linear(din, dh, rngs=rngs)
        self.bn = nnx.BatchNorm(dh, momentum=0.9, rngs=rngs)
        self.drop = nnx.Dropout(0.1, rngs=rngs)
        self.l2 = nnx.Linear(dh, dout, rngs=rngs)
        self.calls = StepCounter(jnp.array(0, jnp.int32))

    def __call__(self, x):
        self.calls[...] += 1                           # modify custom state in place during the forward pass
        x = nnx.relu(self.bn(self.l1(x)))
        return self.l2(self.drop(x))

net = BNNet(16, 32, 4, rngs=nnx.Rngs(params=0, dropout=1))
graphdef, p_state, bn_state, rest = nnx.split(net, nnx.Param, nnx.BatchStat, ...)
print("Param     leaves:", jax.tree.map(jnp.shape, nnx.to_pure_dict(p_state)))
print("BatchStat leaves:", jax.tree.map(jnp.shape, nnx.to_pure_dict(bn_state)))
print("Remaining state:", list(nnx.to_pure_dict(rest).keys()))
assert set(nnx.to_pure_dict(bn_state)["bn"].keys()) == {"mean", "var"}
assert "calls" in nnx.to_pure_dict(rest) and "drop" in nnx.to_pure_dict(rest)
print("BatchStat is a subclass of Variable:", issubclass(nnx.BatchStat, nnx.Variable), "| so is Param:", issubclass(nnx.Param, nnx.Variable))

schedule = optax.warmup_cosine_decay_schedule(init_value=0.0, peak_value=1e-3,
                                              warmup_steps=100, decay_steps=1000, end_value=1e-5)
for t in [0, 50, 100, 550, 1000, 2000]:
    print(f"step {t:5d}: lr = {float(schedule(t)):.3e}")

def cosine_ref(t, W=100, D=1000, peak=1e-3, end=1e-5):
    if t < W:
        return peak * t / W
    t = min(t, D)
    return end + (peak - end) / 2 * (1 + np.cos(np.pi * (t - W) / (D - W)))
for t in [0, 37, 100, 321, 550, 999, 1000, 2000]:
    assert np.isclose(float(schedule(t)), cosine_ref(t), rtol=1e-5, atol=1e-9), t

tx = optax.chain(
    optax.clip_by_global_norm(1.0),                       # clip first
    optax.adamw(learning_rate=schedule, weight_decay=1e-4),
)

# Check the clipping: scale the gradients by hand to global norm 1, feed them to adamw, and the result should match the chain exactly
p0 = {"w": jnp.arange(6.0).reshape(2, 3), "b": jnp.ones(3)}
g0 = {"w": jnp.full((2, 3), 10.0), "b": jnp.full(3, -5.0)}
u_chain, _ = tx.update(g0, tx.init(p0), p0)
gn = optax.global_norm(g0)
g_clipped = jax.tree.map(lambda g: g / gn, g0)
aw = optax.adamw(learning_rate=schedule, weight_decay=1e-4)
u_manual, _ = aw.update(g_clipped, aw.init(p0), p0)
print(f"\nGlobal norm before clipping {float(gn):.2f} → after clipping {float(optax.global_norm(g_clipped)):.2f}")
assert max_diff(u_chain, u_manual) < 1e-9
print("chain(clip, adamw) matches manual clipping + adamw ✓")

# Train BNNet with nnx.Optimizer: in one step the parameters, BatchStat, the custom counter and the optimizer step all get updated
rng = np.random.default_rng(0)
Xb = jnp.asarray(rng.normal(size=(64, 16)), jnp.float32)
Yb = jnp.asarray(rng.normal(size=(64, 4)), jnp.float32)

net = BNNet(16, 32, 4, rngs=nnx.Rngs(params=0, dropout=1))
opt = nnx.Optimizer(net, tx, wrt=nnx.Param)

@nnx.jit
def train_step(model, optimizer, x, y):
    def loss_fn(model):
        return jnp.mean((model(x) - y) ** 2)
    loss, grads = nnx.value_and_grad(loss_fn)(model)   # differentiate only w.r.t. nnx.Param (the default filter)
    optimizer.update(model, grads)                      # update parameters and optimizer state in place
    return loss

bn_before = nnx.to_pure_dict(nnx.state(net, nnx.BatchStat))
w_before = net.l1.kernel[...]
losses = [float(train_step(net, opt, Xb, Yb)) for _ in range(5)]
print("loss:", np.round(losses, 4))
print("optimizer.step =", int(opt.step[...]), "| custom counter calls =", int(net.calls[...]))
assert int(opt.step[...]) == 5 and int(net.calls[...]) == 5
assert max_diff(bn_before, nnx.to_pure_dict(nnx.state(net, nnx.BatchStat))) > 0     # BN statistics were updated
assert max_diff(w_before, net.l1.kernel[...]) > 0                                   # parameters were updated

net.eval()                        # inference mode: BN uses running statistics, Dropout is off
y_eval1, y_eval2 = net(Xb), net(Xb)
assert max_diff(y_eval1, y_eval2) == 0.0
net.train()
print("The forward pass is deterministic in eval mode ✓")

# ❌ Old style (outdated): pmap + one key per device; with jit + sharding's global view you only need one key (see the next cell)
try:
    num_devices = jax.local_device_count()
    key = jax.random.PRNGKey(42)
    device_keys = jax.random.split(key, num_devices)       # split once on the host

    @partial(jax.pmap, axis_name='devices')
    def dropout_pmap(local_key, x):
        dropout_key, new_local_key = jax.random.split(local_key)
        keep = jax.random.bernoulli(dropout_key, 0.5, x.shape)
        return jnp.where(keep, x / 0.5, 0.0), new_local_key

    xs = jnp.ones((num_devices, 4, 8))
    masked, device_keys = dropout_pmap(device_keys, xs)
    same = [bool(jnp.all(masked[0] == masked[i])) for i in range(1, num_devices)]
    print("Measured: pmap still runs; number of devices whose mask equals device 0's =", sum(same))
except Exception as e:
    print("Measured error:", type(e).__name__, e)

# ✅ Recommended: jit + NamedSharding is a "global view": generate random numbers for the whole global array, and each device automatically gets a different slice
mesh_dp = jax.make_mesh((8,), ("data",))
x_global = jnp.ones((32, 8))
key = jax.random.key(42)

@jax.jit
def dropout_global(key, x):
    keep = jax.random.bernoulli(key, 0.5, x.shape)
    return jnp.where(keep, x / 0.5, 0.0)

with jax.set_mesh(mesh_dp):
    x_sharded = jax.device_put(x_global, P("data"))
    y_sharded = dropout_global(key, x_sharded)
    print("Input:", jax.typeof(x_sharded), "→ output:", jax.typeof(y_sharded))
y_single = dropout_global(key, x_global)

# jax_threefry_partitionable defaults to True: sharded and single-device computation produce exactly the same random mask
print("jax_threefry_partitionable =", jax.config.jax_threefry_partitionable)
assert jax.config.jax_threefry_partitionable
assert bool(jnp.all(y_sharded == y_single))
blocks = np.asarray(y_sharded).reshape(8, 4, 8)
assert not all(np.array_equal(blocks[0], blocks[i]) for i in range(1, 8))
print("The sharded mask matches the single-device mask element for element, and each device's slice is different ✓")

# ❌ Wrong (outdated): jax.device_put_replicated was removed in JAX 0.12
try:
    restored_state = {"w": jnp.ones(3)}
    state = jax.device_put_replicated(restored_state, jax.local_devices())
    print("Measured: it runs")
except AttributeError as e:
    print("Measured error: AttributeError:", e)

# ✅ Recommended: save / restore NNX state with Orbax; on restore, data goes straight to the target devices according to the "abstract state" (shape + dtype + sharding)
import orbax.checkpoint as ocp

ckpt_dir = tempfile.mkdtemp(prefix="adv06_ckpt_")
net_src = BNNet(16, 32, 4, rngs=nnx.Rngs(params=0, dropout=1))
opt_src = nnx.Optimizer(net_src, tx, wrt=nnx.Param)
train_step(net_src, opt_src, Xb, Yb)

def pack(model, optimizer):
    """Training state to save: parameters, BN statistics, optimizer (including the step). All pure PyTrees."""
    return {"params": nnx.state(model, nnx.Param),
            "batch_stats": nnx.state(model, nnx.BatchStat),
            "optimizer": nnx.state(optimizer)}

checkpointer = ocp.StandardCheckpointer()
checkpointer.save(os.path.join(ckpt_dir, "step_1"), pack(net_src, opt_src))
checkpointer.wait_until_finished()                 # saving is asynchronous; wait for it to finish before reading back

# Restore: first build an "abstract" model (shapes only, no memory allocated) and use its state as the restore template
abs_net = nnx.eval_shape(lambda: BNNet(16, 32, 4, rngs=nnx.Rngs(params=0, dropout=1)))
abs_opt = nnx.eval_shape(lambda: nnx.Optimizer(abs_net, tx, wrt=nnx.Param))
restored = checkpointer.restore(os.path.join(ckpt_dir, "step_1"), pack(abs_net, abs_opt))

net_new = BNNet(16, 32, 4, rngs=nnx.Rngs(params=123, dropout=1))   # arbitrary initialization
opt_new = nnx.Optimizer(net_new, tx, wrt=nnx.Param)
nnx.update(net_new, restored["params"], restored["batch_stats"])   # write the restored state back
nnx.update(opt_new, restored["optimizer"])
print("optimizer.step after restore =", int(opt_new.step[...]))
assert int(opt_new.step[...]) == 1
assert max_diff(nnx.state(net_new, nnx.Param), nnx.state(net_src, nnx.Param)) == 0.0
assert max_diff(nnx.state(net_new, nnx.BatchStat), nnx.state(net_src, nnx.BatchStat)) == 0.0
assert max_diff(nnx.state(opt_new), nnx.state(opt_src)) == 0.0
print("Orbax save → restore: parameters, BN statistics and optimizer state are bit-identical ✓")

import flax.linen as nn
from flax.training import train_state

class LinenCNN(nn.Module):               # a minimal model added here to serve as `model`
    @nn.compact
    def __call__(self, x):
        x = x.reshape((x.shape[0], -1))
        x = nn.relu(nn.Dense(32)(x))
        return nn.Dense(10)(x)

model = LinenCNN()
key = jax.random.key(0)
lparams = model.init(key, jnp.ones([1, 28, 28, 1]))['params']
tx_adam = optax.adam(learning_rate=1e-3)
state = train_state.TrainState.create(apply_fn=model.apply, params=lparams, tx=tx_adam)
num_devices = jax.local_device_count()

# ❌ Wrong (outdated): jax.device_put_replicate is misspelled, and the correctly spelled jax.device_put_replicated was also removed in 0.12
for fn_name in ["device_put_replicate", "device_put_replicated"]:
    try:
        replicated_state = getattr(jax, fn_name)(state, jax.local_devices())
        print(fn_name, "→ it runs")
    except AttributeError as e:
        print(f"jax.{fn_name} → AttributeError: {str(e)[:110]}")

# ✅ Recommended (if you still want pmap): following the official migration guide, use jax.device_put to place a "stacked, split by device" array
def replicate_for_pmap(tree, devices):
    mesh = jax.sharding.Mesh(np.array(devices), ("d",))
    n = len(devices)
    return jax.tree.map(
        lambda x: jax.device_put(jnp.broadcast_to(x, (n,) + jnp.shape(x)), NamedSharding(mesh, P("d"))),
        tree)

replicated_state = replicate_for_pmap(state, jax.local_devices())
print(f"State replicated across {num_devices} devices.")
print("kernel shape after replication:", replicated_state.params['Dense_0']['kernel'].shape)
print("original kernel shape:         ", state.params['Dense_0']['kernel'].shape)
assert replicated_state.params['Dense_0']['kernel'].shape == (num_devices,) + state.params['Dense_0']['kernel'].shape

# Course steps 2 and 3: shard the data + pmapped training step
global_batch_size = 64
local_batch_size = global_batch_size // num_devices

def shard_batch(batch):
    """Reshape (B, ...) into (num_devices, B/num_devices, ...)."""
    return jax.tree.map(lambda x: x.reshape((num_devices, local_batch_size) + x.shape[1:]), batch)

rng = np.random.default_rng(0)
global_images = jnp.asarray(rng.normal(size=(global_batch_size, 28, 28, 1)), jnp.float32)
global_labels = jnp.asarray(rng.integers(0, 10, global_batch_size), jnp.int32)
batch = {'image': global_images, 'label': global_labels}
sharded_batch = shard_batch(batch)
print("Global image shape:", global_images.shape, "→ sharded image shape:", sharded_batch['image'].shape)

def compute_loss(params, batch, apply_fn):
    logits = apply_fn({'params': params}, batch['image'])
    one_hot = jax.nn.one_hot(batch['label'], num_classes=10)
    return jnp.mean(-jnp.sum(one_hot * jax.nn.log_softmax(logits), axis=-1))

def train_step(state, batch):
    loss, grads = jax.value_and_grad(compute_loss)(state.params, batch, state.apply_fn)
    grads = jax.lax.pmean(grads, axis_name='batch_axis')      # key step: average gradients across devices
    return state.apply_gradients(grads=grads), loss

# ❌ Old style (outdated): pmap still works in 0.12, but the official advice is to migrate to jit + sharding (or shard_map)
try:
    p_train_step = jax.pmap(train_step, axis_name='batch_axis')
    pmap_state, local_losses = p_train_step(replicated_state, sharded_batch)
    print("Measured: pmap runs; local_losses shape", local_losses.shape, "| mean loss", f"{float(local_losses.mean()):.4f}")
    print("Sharding type of the pmap output:", type(local_losses.sharding).__name__, local_losses.sharding.spec)
except Exception as e:
    print("Measured error:", type(e).__name__, e)

# ✅ Recommended: NNX + Optax + jit + NamedSharding data parallelism (no pmean / manual reshape / state replication needed)
class NNXMLP(nnx.Module):
    def __init__(self, *, rngs):
        self.l1 = nnx.Linear(784, 32, rngs=rngs)
        self.l2 = nnx.Linear(32, 10, rngs=rngs)
    def __call__(self, x):
        return self.l2(nnx.relu(self.l1(x.reshape(x.shape[0], -1))))

@nnx.jit
def nnx_train_step(model, optimizer, x, y):
    def loss_fn(model):
        return optax.softmax_cross_entropy_with_integer_labels(model(x), y).mean()
    loss, grads = nnx.value_and_grad(loss_fn)(model)
    optimizer.update(model, grads)
    return loss

mesh_dp = jax.make_mesh((8,), ("data",))
print("mesh:", mesh_dp)
with jax.set_mesh(mesh_dp):
    dp_model = NNXMLP(rngs=nnx.Rngs(0))
    dp_opt = nnx.Optimizer(dp_model, optax.adam(1e-3), wrt=nnx.Param)
    x_dp = jax.device_put(global_images, P("data"))          # shard the batch dimension
    y_dp = jax.device_put(global_labels, P("data"))
    print("x:", jax.typeof(x_dp), "| kernel:", jax.typeof(dp_model.l1.kernel[...]))
    dp_losses = [float(nnx_train_step(dp_model, dp_opt, x_dp, y_dp)) for _ in range(3)]
    jax.debug.visualize_array_sharding(x_dp.reshape(64, -1)[:, :8])

# Single-device reference: same initialization, same data
ref_model = NNXMLP(rngs=nnx.Rngs(0))
ref_opt = nnx.Optimizer(ref_model, optax.adam(1e-3), wrt=nnx.Param)
ref_losses = [float(nnx_train_step(ref_model, ref_opt, global_images, global_labels)) for _ in range(3)]
print("Data-parallel loss:", np.round(dp_losses, 6))
print("Single-device loss:", np.round(ref_losses, 6))
d = max_diff(nnx.state(dp_model, nnx.Param), nnx.state(ref_model, nnx.Param))
print(f"Max parameter difference after 3 steps: {d:.2e}")
assert np.allclose(dp_losses, ref_losses, rtol=1e-5) and d < 1e-5

# Take a look at the communication the compiler inserted automatically
with jax.set_mesh(mesh_dp):
    hlo = nnx_train_step.lower(dp_model, dp_opt, x_dp, y_dp).compile().as_text()
print("Automatically inserted collectives:", collectives(hlo))
assert collectives(hlo)["all-reduce"] >= 1

# The pmap version also matches the modern style (the course's Linen model vs single-device jit)
@jax.jit
def single_step(state, batch):
    loss, g = jax.value_and_grad(compute_loss)(state.params, batch, state.apply_fn)
    return state.apply_gradients(grads=g), loss

s1, l1 = single_step(state, batch)
d = max_diff(jax.tree.map(lambda x: x[0], pmap_state.params), s1.params)
print(f"pmap replica 0 vs single device: parameter difference {d:.2e}; mean loss {float(local_losses.mean()):.6f} vs {float(l1):.6f}")
assert d < 1e-6 and np.isclose(float(local_losses.mean()), float(l1), rtol=1e-6)
# All replicas are exactly identical
assert all(max_diff(jax.tree.map(lambda x: x[i], pmap_state.params),
                    jax.tree.map(lambda x: x[0], pmap_state.params)) == 0 for i in range(num_devices))

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

class ShardedMLP(nnx.Module):
    def __init__(self, din, dh, dout, *, rngs):
        init = nnx.initializers.lecun_normal()
        # First layer is split by columns (output dim → model axis), second layer by rows (input dim → model axis)
        self.up = nnx.Linear(din, dh, use_bias=False,
                             kernel_init=nnx.with_partitioning(init, (None, "model")), rngs=rngs)
        self.down = nnx.Linear(dh, dout, use_bias=False,
                               kernel_init=nnx.with_partitioning(init, ("model", None)), rngs=rngs)
    def __call__(self, x):
        h = nnx.relu(self.up(x))
        # In Explicit mode, a matmul whose contracting dim is split needs an explicit output sharding (here: summed, then sharded by batch)
        return self.down(h, out_sharding=P("data", None))

try:
    ShardedMLP(64, 256, 64, rngs=nnx.Rngs(0))
except ValueError as e:
    print("No mesh context → ValueError:", str(e)[:100], "...")

with jax.set_mesh(mesh2d):
    smodel = ShardedMLP(64, 256, 64, rngs=nnx.Rngs(0))
    sopt = nnx.Optimizer(smodel, optax.adam(1e-3), wrt=nnx.Param)
    print("Sharding annotations:", jax.tree.map(str, nnx.to_pure_dict(nnx.get_partition_spec(nnx.state(smodel)))))
    print("up.kernel   :", jax.typeof(smodel.up.kernel[...]))
    print("down.kernel :", jax.typeof(smodel.down.kernel[...]))
    mu_up = sopt.opt_state[0].mu["up"]["kernel"][...]
    print("Adam mu(up) :", jax.typeof(mu_up))
    assert mu_up.sharding.spec == smodel.up.kernel[...].sharding.spec == P(None, "model")
    # Each device stores only 1/4 of up.kernel
    k = smodel.up.kernel[...]
    print(f"up.kernel: {k.nbytes} B global, {k.addressable_shards[0].data.nbytes} B per device")
    assert k.addressable_shards[0].data.nbytes * 4 == k.nbytes
    jax.debug.visualize_array_sharding(k)

rng = np.random.default_rng(0)
X_acc = jnp.asarray(rng.normal(size=(64, 16)), jnp.float32)
Y_acc = jnp.asarray(rng.normal(size=(64, 1)), jnp.float32)
acc_params = {"w1": jnp.asarray(rng.normal(size=(16, 32)) * 0.3, jnp.float32), "b1": jnp.zeros(32),
              "w2": jnp.asarray(rng.normal(size=(32, 1)) * 0.3, jnp.float32), "b2": jnp.zeros(1)}

def mlp_apply(p, x):
    return jnp.tanh(x @ p["w1"] + p["b1"]) @ p["w2"] + p["b2"]

def loss_fn(p, x, y):
    return jnp.mean((mlp_apply(p, x) - y) ** 2)     # per-example mean

acc_tx = optax.adam(1e-2)
N_ACC = 4

def micro_batches(x, y, n):
    return [(x[i * len(x) // n:(i + 1) * len(x) // n], y[i * len(y) // n:(i + 1) * len(y) // n]) for i in range(n)]

# ❌ Wrong: apply_update is defined inside the function, so every call creates a new function object → it's retraced/recompiled every step
trace_count = {"n": 0}

@jax.jit
def micro_batch_step(params, x, y):
    return jax.value_and_grad(loss_fn)(params, x, y)

def accumulated_train_step(params, opt_state, data_iterator, accumulation_steps):
    accumulated_grads = jax.tree.map(jnp.zeros_like, params)
    total_loss = 0.0
    for _ in range(accumulation_steps):
        x, y = next(data_iterator)
        loss, grads = micro_batch_step(params, x, y)
        accumulated_grads = jax.tree.map(lambda a, g: a + g, accumulated_grads, grads)
        total_loss += loss
    averaged_grads = jax.tree.map(lambda g: g / accumulation_steps, accumulated_grads)

    @jax.jit
    def apply_update(p, o_state, avg_grads):
        trace_count["n"] += 1                       # only runs during tracing
        updates, new_o_state = acc_tx.update(avg_grads, o_state, p)
        return optax.apply_updates(p, updates), new_o_state

    params, opt_state = apply_update(params, opt_state, averaged_grads)
    return params, opt_state, total_loss / accumulation_steps

try:
    p_course, s_course = acc_params, acc_tx.init(acc_params)
    for step in range(5):
        p_course, s_course, _ = accumulated_train_step(p_course, s_course,
                                                       iter(micro_batches(X_acc, Y_acc, N_ACC)), N_ACC)
    print(f"Measured: in 5 training steps, apply_update was traced {trace_count['n']} times")
except Exception as e:
    print("Measured error:", type(e).__name__, e)
assert trace_count["n"] == 5

# ✅ Recommended: put the whole accumulation loop in one jit and iterate over micro-batches with lax.scan (reshape the data to (N, b, ...) first)
@partial(jax.jit, static_argnames="n_micro")
def accum_grads_scan(params, x, y, n_micro):
    xs = x.reshape(n_micro, -1, *x.shape[1:])
    ys = y.reshape(n_micro, -1, *y.shape[1:])
    def body(carry, mb):
        acc, loss_sum = carry
        loss, g = jax.value_and_grad(loss_fn)(params, *mb)
        return (jax.tree.map(jnp.add, acc, g), loss_sum + loss), None
    init = (jax.tree.map(jnp.zeros_like, params), 0.0)
    (acc, loss_sum), _ = lax.scan(body, init, (xs, ys))
    return jax.tree.map(lambda a: a / n_micro, acc), loss_sum / n_micro

@partial(jax.jit, static_argnames="n_micro", donate_argnums=(0, 1))
def accum_train_step(params, opt_state, x, y, n_micro):
    grads, loss = accum_grads_scan(params, x, y, n_micro)
    updates, opt_state = acc_tx.update(grads, opt_state, params)
    return optax.apply_updates(params, updates), opt_state, loss

g_big = jax.jit(jax.grad(loss_fn))(acc_params, X_acc, Y_acc)      # the whole batch at once
g_acc, _ = accum_grads_scan(acc_params, X_acc, Y_acc, N_ACC)
d = max_diff(g_big, g_acc)
print(f"lax.scan accumulating 4 micro-batches vs one large batch: max gradient difference {d:.2e}")
assert d < 1e-6

# Also matches the step-by-step accumulation above (5 training steps)
p_new, s_new = jax.tree.map(jnp.copy, acc_params), acc_tx.init(acc_params)
for step in range(5):
    p_new, s_new, _ = accum_train_step(p_new, s_new, X_acc, Y_acc, n_micro=N_ACC)
print(f"After 5 steps, scan version vs step-by-step version parameter difference: {max_diff(p_new, p_course):.2e}")
assert max_diff(p_new, p_course) < 1e-5