code/train/trainer.py

# coding=utf-8
# Copyright 2024 The Google Research Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

"""Transformer LM trainer."""

import functools

from clu import periodic_actions
from flax.training import common_utils
from flax.training import train_state
import jax
from jax import numpy as jnp
import numpy as np
import optax
import ml_collections

from train import model


def get_state(config, net, initial_variables):
    """Get the train state given an experiment config, a model and initial variables.

    Args:
        config: A ConfigDict containing the configuration for the experiment.
        net: The model to use for training.
        initial_variables: The initial variables for the model.

    Returns:
        A tuple containing the train state and the learning rate schedule.
    """
    # Learning rate schedule
    lr_scheduler_fn = functools.partial(
        lr_scheduler,
        learning_rate=config.learning_rate,
        warmup_tokens=config.warmup_tokens,
        final_tokens=config.max_steps,
        config=config,
    )
    # Optimizer
    optim_fn = optax.adamw(
        lr_scheduler_fn, weight_decay=config.weight_decay, b1=0.9, b2=0.95
    )
    # Clip the gradients to prevent exploding gradients.
    optimizer = optax.chain(optax.clip_by_global_norm(1), optim_fn)

    # Initialize the train state
    state = train_state.TrainState.create(
        apply_fn=net.apply, params=initial_variables["params"],
        tx=optimizer
        )

    return state, lr_scheduler_fn

def lr_scheduler(
    n_tokens: int, learning_rate: float, warmup_tokens: int, final_tokens: int,
    config: ml_collections.ConfigDict,
) -> float:
    """Learning rate scheduler, adapted from Mikhail Grankin.

    The learning rate schedule is cosine decay with a warmup period.

    The learning rate starts at 0 and linearly increases to the given learning
    rate over the warmup period. After the warmup period, the learning rate
    decays according to a cosine schedule, with the given learning rate as the
    maximum value.

    Args:
        n_tokens: The number of tokens processed so far.
        learning_rate: The initial learning rate.
        warmup_tokens: The number of tokens to warm up over.
        final_tokens: The total number of tokens to process.
        config: A ConfigDict containing the configuration for the learning rate
            schedule.

    Returns:
        The learning rate at the given point in the schedule.
    """
    # Decay the learning rate based on our progress.
    progress = (n_tokens - warmup_tokens) / max(
        1, final_tokens - warmup_tokens,
    )
    lr_mult = jnp.where(
        n_tokens < warmup_tokens,
        # Linear warmup.
        n_tokens / jnp.fmax(1, warmup_tokens),
        # Cosine learning rate decay.
        jnp.fmax(config.end_lr_factor, 0.5 * (1.0 + jnp.cos(np.pi * progress))),
    )
    return learning_rate * lr_mult


def get_metrics_report_progress(config, workdir, writer):
    """
    Get the metrics for reporting progress during training.

    Args:
        config: The configuration for the experiment.
        workdir: The directory for storing the logs.
        writer: The writer object for recording the metrics.

    Returns:
        hooks: List of hooks for tracking progress.
        report_progress: Object for reporting progress.
        train_metrics: List of training metrics.
    """
    hooks = []

    # Initialize the report progress object
    report_progress = periodic_actions.ReportProgress(
        num_train_steps=config.max_steps, writer=writer)

    # Add metrics for profiling if the process index is 0
    if jax.process_index() == 0:
        hooks += [report_progress,
                  periodic_actions.Profile(logdir=workdir, num_profile_steps=5)]
    
    # Initialize the list of training metrics
    train_metrics = []
    
    return hooks, report_progress, train_metrics


def get_input_start_index(batch, config):
    inputs = jax.tree_util.tree_map(np.asarray, batch[0])
    puzzles = jax.tree_util.tree_map(np.asarray, batch[1])
    start_index = jax.tree_util.tree_map(np.asarray, batch[2])
    levels = jax.tree_util.tree_map(np.asarray, batch[3])
    cand_targets = jax.tree_util.tree_map(np.asarray, batch[4])
    return inputs, puzzles, start_index, levels, cand_targets


def train_one_step(p_train_step, config, state, step, dropout_rngs, train_data_iter):
    """
    Single step of the training loop.

    Args:
        p_train_step: The training step function.
        config: The experiment configuration.
        state: The train state.
        step: The step number.
        dropout_rngs: The dropout random number generator.
        train_data_iter: The iterator for the train data.

    Returns:
        The updated train state and train metrics.
    """
    with jax.profiler.StepTraceAnnotation("train", step_num=step):
        # Get the next batch from the iterator
        batch = next(train_data_iter)
        # Extract the inputs, start index and difficulty level from the batch
        inputs, _, start_index, levels, cand_targets = get_input_start_index(batch, config)
        # Shard across the devices
        inputs = common_utils.shard(jax.tree_util.tree_map(np.asarray, inputs))
        start_index = common_utils.shard(jax.tree_util.tree_map(np.asarray, start_index))
        levels = common_utils.shard(jax.tree_util.tree_map(np.asarray, levels))
        # Explicit reshape (not common_utils.shard) so a zero-width slot dim
        # (K=0 no-latent baseline: cand_targets is (bs, 0, 81)) shards without
        # the ambiguous -1 inference. Identical to shard() when K>0.
        cand_targets = np.asarray(cand_targets)
        _nd = jax.local_device_count()
        cand_targets = cand_targets.reshape(
            (_nd, cand_targets.shape[0] // _nd) + cand_targets.shape[1:])

        # Run the training step
        state, metrics, _ = p_train_step(
            state, inputs, start_index, levels, cand_targets, dropout_rng=dropout_rngs
        )

    return state, metrics


def _shard_train_batch(batch, config):
    """Unpack a dataset batch and shard it the way train_one_step does."""
    inputs, _, start_index, levels, cand_targets = get_input_start_index(
        batch, config)
    inputs = common_utils.shard(jax.tree_util.tree_map(np.asarray, inputs))
    start_index = common_utils.shard(
        jax.tree_util.tree_map(np.asarray, start_index))
    levels = common_utils.shard(jax.tree_util.tree_map(np.asarray, levels))
    cand_targets = np.asarray(cand_targets)
    nd = jax.local_device_count()
    cand_targets = cand_targets.reshape(
        (nd, cand_targets.shape[0] // nd) + cand_targets.shape[1:])
    return inputs, start_index, levels, cand_targets


def apply_grads(state, grads):
    """Replicated optimizer step. Used after averaging per-depth grads."""
    return state.apply_gradients(grads=grads)


def train_depth_mix_step(grad_step_for, p_apply, config, state, step,
                         dropout_rngs, mix_ds, n_depths):
    """One optimizer step whose CE is the mean over depths 1..n_depths.

    Depth d is run with d latent passes on a batch whose value tokens are
    stage-(d-1) instances, drawn from `mix_ds` (a SudokuDataset, not tf.data).
    Gradients are summed and divided by n_depths, then applied once, so the
    update is joint rather than a serial overwrite.
    """
    n_depths = int(max(n_depths, 1))
    bs = int(config.minibatch_size)
    acc = None
    metrics_list = []
    with jax.profiler.StepTraceAnnotation("train", step_num=step):
        for d in range(1, n_depths + 1):
            batch = mix_ds.sample_instance_batch(bs, d - 1)
            inputs, start_index, levels, cand_targets = _shard_train_batch(
                batch, config)
            grads, metrics, _ = grad_step_for(d)(
                state, inputs, start_index, levels, cand_targets,
                dropout_rng=dropout_rngs)
            acc = (grads if acc is None
                   else jax.tree_util.tree_map(lambda a, b: a + b, acc, grads))
            metrics_list.append(metrics)
        acc = jax.tree_util.tree_map(lambda a: a / n_depths, acc)
        state = p_apply(state, acc)

    def _mean_metric(key):
        vals = [m[key] for m in metrics_list if key in m]
        if not vals:
            return None
        return sum(vals) / float(len(vals))

    metrics = dict(metrics_list[-1])
    for key in ("loss", "ce_loss", "aux_loss", "leak_loss", "latent_sigma"):
        v = _mean_metric(key)
        if v is not None:
            metrics[key] = v
    return state, metrics


def build_latent_state(inputs, start_index, levels, config, num_passes,
                       apply_fn, rngs=None, trunc_passes=0):
    """Build the continuous latent thoughts (Coconut/ATC-style recurrence).

    Pass 0 seeds z_1 from the last-layer hidden at the last clue token; pass j
    reads the hidden at latent slot j to produce z_{j+1}. Slot j is active for
    an example only when j < k, with k = clip(level - 2, 1, K) (difficulty-
    matched latent budget). By default gradients flow through all passes
    (full BPTT). If trunc_passes=T>0, the forward chain is unchanged but
    thoughts older than the last T CoTs are stop-gradient'd, so the loss
    only backprops through those last T recurrent steps.

    Returns (latent_vals, slot_pos, active_full).
    """
    num_slots = int(config.num_latent_slots)
    bs = inputs.shape[0]
    bidx = jnp.arange(bs)

    si3 = 3 * start_index.reshape(-1)                      # (bs,)
    slot_pos = si3[:, None] + jnp.arange(num_slots)[None, :]   # (bs, K)
    if getattr(config, "cand_slot_mode", "level") == "depth":
        # Every slot the recurrence fills is active, uniformly over the batch.
        k = jnp.full((bs,), max(min(num_passes, num_slots), 1), dtype=jnp.int32)
    else:
        k = jnp.clip(levels.reshape(-1) - 2, 1, num_slots)  # (bs,)
    active_full = jnp.arange(num_slots)[None, :] < k[:, None]  # (bs, K)

    latent_vals = jnp.zeros((bs, num_slots, config.emb_dim), dtype=config.dtype)

    for j in range(num_passes):
        act_j = active_full & (jnp.arange(num_slots)[None, :] < j)
        _, hidden, _ = apply_fn(inputs, latent_vals, slot_pos, act_j, rngs)
        src = si3 - 1 + j   # j=0: last clue token; j>0: latent slot j-1
        z = hidden[bidx, src]
        # Truncated BPTT: keep the full forward chain, but detach CoTs
        # before the last T so earlier thoughts are treated as constants.
        if trunc_passes > 0 and j < num_passes - trunc_passes:
            z = jax.lax.stop_gradient(z)
        latent_vals = latent_vals.at[:, j].set(z.astype(config.dtype))

    return latent_vals, slot_pos, active_full


def build_latent_state_noisy(inputs, start_index, levels, config, num_passes,
                             apply_fn, sample_slot, rng, trunc_passes=0):
    """One sample: noise h1, pass, noise h2, pass, ..., then readout."""
    num_slots = int(config.num_latent_slots)
    bs = inputs.shape[0]
    bidx = jnp.arange(bs)
    si3 = 3 * start_index.reshape(-1)
    slot_pos = si3[:, None] + jnp.arange(num_slots)[None, :]
    if getattr(config, "cand_slot_mode", "level") == "depth":
        k = jnp.full((bs,), max(min(num_passes, num_slots), 1), dtype=jnp.int32)
    else:
        k = jnp.clip(levels.reshape(-1) - 2, 1, num_slots)
    active_full = jnp.arange(num_slots)[None, :] < k[:, None]
    latent_vals = jnp.zeros((bs, num_slots, config.emb_dim), dtype=config.dtype)
    keys = jax.random.split(rng, num_passes)
    last_sigma = jnp.zeros((), jnp.float32)

    for j in range(num_passes):
        act_j = active_full & (jnp.arange(num_slots)[None, :] < j)
        drop_k, noise_k = jax.random.split(keys[j])
        _, hidden, _ = apply_fn(
            inputs, latent_vals, slot_pos, act_j, {"dropout": drop_k})
        # Same read position as the deterministic recurrence and the gate:
        # j=0 is the last clue token, j>0 is latent slot j-1's hidden.
        z = hidden[bidx, si3 - 1 + j]
        if trunc_passes > 0 and j < num_passes - trunc_passes:
            z = jax.lax.stop_gradient(z)
        mu, log_sigma = sample_slot(z, j)
        log_sigma = jnp.clip(log_sigma, -12.0, 2.0)
        eps = jax.random.normal(noise_k, z.shape, jnp.float32)
        z = (mu + jnp.exp(log_sigma) * eps).astype(config.dtype)
        last_sigma = jnp.exp(log_sigma).mean()
        latent_vals = latent_vals.at[:, j].set(z)

    return latent_vals, slot_pos, active_full, last_sigma


def train_step(state, batch, start_index, levels, cand_targets, config,
               hyperparams, learning_rate_fn, num_passes, dropout_rng=None,
               backtrack=False, inst_slot=-1, apply_update=True):
    """One step of the training loop.

    Args:
        state: Train state.
        batch: Input batch (bs, 3*81 + K) with latent placeholder slots.
        start_index: Number of clue cells per example.
        levels: Puzzle difficulty level (3..8) per example.
        config: Model config.
        hyperparams: Hyperparameter dictionary.
        learning_rate_fn: Learning rate function.
        num_passes: Number of latent recurrence passes (= curriculum stage).
        dropout_rng: RNG used for dropout.
        inst_slot: Candidate-mask slot matching the stage whose instances supply
            this batch's value tokens (= curriculum stage - 1). Only the leak
            penalty reads it. -1 derives it from num_passes.

    Returns:
        A new train state, train metrics, and computed model predictions.
    """
    num_slots = int(config.num_latent_slots)
    # Extract inputs and labels from the batch
    inputs = batch[:, :-1]
    label = batch[:, 1:]

    # Update dropout_rng
    dropout_rng = jax.random.fold_in(dropout_rng, state.step)
    # Stochastic latents draw from their own stream. Only split when the arm is
    # on, so the deterministic control keeps its exact dropout sequence.
    n_mc = int(getattr(hyperparams, "latent_mc_samples", 0))
    use_noise = bool(getattr(config, "latent_noise", False)) and n_mc > 0
    noise_rng = None
    if use_noise:
        dropout_rng, noise_rng = jax.random.split(dropout_rng)
    mc_objective = str(getattr(hyperparams, "latent_mc_objective", "mean"))
    dropout_rng_dict = {"dropout": dropout_rng}

    def loss_fn(params):
        """Compute the loss function."""
        net = model.TransformerLMHeadModel(config)

        def apply_fn(x, lv, lp, la, rngs):
            return net.apply({"params": params}, x, latent_values=lv,
                             latent_positions=lp, latent_active=la,
                             rngs=rngs)

        # M independent samples of (noise h1, pass, noise h2, ..., readout).
        # Mean CE.
        sigma_mean = jnp.zeros((), jnp.float32)
        if use_noise:
            def _sample_slot(z, slot):
                return net.apply(
                    {"params": params}, z, slot,
                    method=model.TransformerLMHeadModel.latent_noise_dist_slot)

            def _one_draw(key):
                build_key, read_key = jax.random.split(key)
                def apply_draw(x, lv, lp, la, rngs):
                    return apply_fn(x, lv, lp, la, rngs)
                latent_vals, slot_pos, active_full, sig = (
                    build_latent_state_noisy(
                        inputs, start_index, levels, config, num_passes,
                        apply_draw, _sample_slot, build_key,
                        trunc_passes=int(getattr(hyperparams, "bptt_trunc", 0))))
                lg, _, cl = apply_draw(
                    inputs, latent_vals, slot_pos, active_full,
                    {"dropout": read_key})
                return lg, cl, sig

            pred_logits_mc, cand_logits_mc, sig_mc = jax.vmap(_one_draw)(
                jax.random.split(noise_rng, n_mc))
            pred_logits = pred_logits_mc.mean(axis=0)
            cand_logits = cand_logits_mc.mean(axis=0)
            sigma_mean = sig_mc.mean()
        else:
            latent_vals, slot_pos, active_full = build_latent_state(
                inputs, start_index, levels, config, num_passes, apply_fn,
                rngs=dropout_rng_dict,
                trunc_passes=int(getattr(hyperparams, "bptt_trunc", 0)))
            pred_logits, _, cand_logits = apply_fn(
                inputs, latent_vals, slot_pos, active_full, dropout_rng_dict)
            pred_logits_mc = pred_logits[None]

        label_one_hot = jax.nn.one_hot(label, num_classes=config.vocab_size)

        # The variables label_one_hot and pred_logits both are 3-dimensional tensors with
        # first axis corresponding to batch size, second correspondingn to sequence length 
        # and third corresponding to the row/column/value at a particular cell
        assert label_one_hot.shape == pred_logits.shape, ("one hot label shape",
                                                        label_one_hot.shape,
                                                        label.shape,
                                                        pred_logits.shape)
        
        # Calculate the cross-entropy loss along the last axis, per MC draw.
        # ce_loss is (M, bs, L); M is 1 on the deterministic path, so the
        # arithmetic below reduces to the original single-forward loss exactly.
        pred_logits_sol = pred_logits[:, :, :]
        label_one_hot_sol = label_one_hot[:, :, :]

        ce_loss = optax.softmax_cross_entropy(
            logits=pred_logits_mc, labels=label_one_hot_sol[None]
        )

        # Apply masking to the loss: supervise only the solution region,
        # which now starts K latent slots after the clue block.
        pos = jnp.arange(label.shape[1])[None, :]                   # (1, L)
        base = 3 * jnp.asarray(start_index).reshape(-1, 1) + num_slots
        mask = pos >= base

        # Split the supervised tokens by role. The solution region is a stream of
        # (row, col, value) triples, so the value token of empty cell t sits at
        # base+3t+2; ce_loss[i] scores the token at i+1, hence the shift below.
        # Two of every three supervised tokens are locations, and they demand the
        # solver's particular linearization of a wave whose cells are actually
        # order-interchangeable -- pressure the promotion gate (a teacher-forced,
        # value-only CE(q||p)) never measures. Down-weighting them spends the
        # capacity on the value distribution instead.
        is_value = ((pos + 1 - base) % 3) == 2
        loc_weight = float(getattr(hyperparams, "loc_loss_weight", 1.0))
        value_weight = float(getattr(hyperparams, "value_loss_weight", 1.0))
        token_w = jnp.where(is_value, value_weight, loc_weight) * mask

        # Per-token mean (not per-example sum): keeps the LM CE on the same
        # O(1) scale as the per-digit BCE below, so aux_cand_weight~1 actually
        # balances the two instead of the candidate signal being swamped.
        # The denominator stays the *unweighted* supervised count so that
        # re-weighting moves the location gradient without silently rescaling the
        # value gradient too, and weights of 1.0 reproduce the plain mean exactly.
        ce_denom = jnp.maximum(mask.sum(), 1.0)
        # Per-draw scalar first, then reduce over the draws. "mean" is the
        # Monte-Carlo estimate of E_z[CE]; "min" charges only the best draw,
        # which is what lets separate draws specialize instead of all being
        # pulled to the same target.
        per_draw_ce = (ce_loss * token_w[None]).sum(axis=(1, 2)) / ce_denom
        if use_noise and mc_objective == "min":
            avg_ce_loss = per_draw_ce.min()
        else:
            avg_ce_loss = per_draw_ce.mean()

        # ---- Leak penalty: probability mass placed outside the candidate set --
        # The gate is CE(q||p) <= log|S|, and CE(q||p) = -log(in-set mass)
        # + H(q) + KL(q||p_renormalized). Measured on a stage-3 checkpoint the
        # leak term is what grows with curriculum depth (0.03 nats at stage 3,
        # 0.07 at stage 4, 0.13 at stage 5) while the within-set shape error
        # stays flat, so it accounts for the whole deepening deficit. Nothing
        # else in this objective penalizes it: the LM CE only ever rewards the
        # one sampled digit, and the candidate head that used to supply the set
        # boundary is off (aux_cand_weight=0) on the instance arm. This term
        # minimizes exactly the measured quantity, -log(sum of p over S).
        leak_weight = float(getattr(hyperparams, "leak_weight", 0.0))
        if leak_weight > 0.0:
            slot = inst_slot if inst_slot >= 0 else num_passes - 1
            slot = int(min(max(slot, 0), num_slots - 1))
            # At a value position i the row/col tokens sit at i-1 and i, so the
            # cell this position is answering is inputs[i-1]*9 + inputs[i].
            prev_tok = jnp.concatenate(
                [jnp.zeros((inputs.shape[0], 1), inputs.dtype),
                 inputs[:, :-1]], axis=1)
            cell = jnp.clip(prev_tok.astype(jnp.int32) * 9
                            + inputs.astype(jnp.int32), 0, 80)
            bits_at = jnp.take_along_axis(
                cand_targets[:, slot, :].astype(jnp.int32), cell, axis=1)
            dv = jnp.arange(config.vocab_size)
            # Token ids are the digits themselves, so digit d lives at index d.
            in_set = (
                ((bits_at[:, :, None] >> jnp.clip(dv - 1, 0, 8)[None, None, :])
                 & 1) > 0) & ((dv >= 1) & (dv <= 9))[None, None, :]
            # Shifted softmax in fp32: p <= 1 and Z >= 1, so both logs are safe
            # without a logsumexp import.
            lg = pred_logits_sol.astype(jnp.float32)
            pe = jnp.exp(lg - lg.max(axis=-1, keepdims=True))
            z_all = pe.sum(axis=-1)
            z_in = (pe * in_set).sum(axis=-1)
            leak = jnp.log(z_all) - jnp.log(jnp.maximum(z_in, 1e-30))
            # Only real empty cells: clue cells are sentinel-zeroed in the masks.
            leak_m = mask & is_value & (bits_at > 0)
            leak_loss = ((leak * leak_m).sum()
                         / jnp.maximum(leak_m.sum(), 1.0))
        else:
            leak_loss = jnp.zeros((), dtype=jnp.float32)

        # ---- Auxiliary multi-candidate BCE loss on the latent slots ----
        # cand_targets: (bs, K, 81) int bitmask (bit d-1 set <=> digit d is a
        # candidate at that slot's stage; slot j already aligned to the puzzle's
        # k=level-2 budget in the data pipeline). Expand to (bs,K,81,9) multi-hot.
        aux_weight = float(getattr(hyperparams, "aux_cand_weight", 1.0))
        if aux_weight == 0.0:
            # Instance arm: the superposition lives in the varying targets, not
            # in a set head. Drop the whole BCE graph so nothing but the LM CE
            # shapes the latents.
            zero = jnp.zeros((), dtype=pred_logits.dtype)
            total = avg_ce_loss + leak_weight * leak_loss
            return total, (pred_logits, avg_ce_loss, zero, leak_loss,
                           sigma_mean)

        bits = jnp.arange(9)
        cand_multi_hot = ((cand_targets[..., None].astype(jnp.int32)
                           >> bits) & 1).astype(cand_logits.dtype)  # (bs,K,81,9)
        # Positive-weighted BCE. Candidate masks are sparse (~1-3 of 9 digits
        # "on"), so plain BCE collapses to predicting all-zeros. Up-weighting
        # the positive (candidate-present) term by pos_weight counteracts the
        # imbalance and forces the head to actually predict the candidate set.
        pos_weight = float(getattr(hyperparams, "aux_pos_weight", 5.0))
        log_p = jax.nn.log_sigmoid(cand_logits)                     # log sigmoid(x)
        log_1mp = jax.nn.log_sigmoid(-cand_logits)                  # log(1 - sigmoid(x))
        bce = -(pos_weight * cand_multi_hot * log_p
                + (1.0 - cand_multi_hot) * log_1mp)                 # (bs,K,81,9)
        # Only the empty cells are supervised: clue cells were sentinel-zeroed
        # in the data pipeline, so any cell whose target row is all-zero is a
        # clue and must not contribute. This makes the effective target
        # (#empty cells) x 9 per puzzle rather than 81 x 9.
        cell_mask = (cand_targets > 0).astype(cand_logits.dtype)    # (bs, K, 81)
        # Delta weighting: consecutive stages are near-duplicates (at K=12 only
        # ~14% of cells change per stage, so ~97% of the target bits are copies
        # of the previous slot). Weighting the unchanged cells below 1 stops the
        # objective from being satisfied by echoing slot j-1. Slot 0 has no
        # predecessor, so it is fully weighted.
        delta_bg = float(getattr(hyperparams, "aux_delta_bg", 1.0))
        if delta_bg != 1.0:
            changed = jnp.concatenate(
                [jnp.ones_like(cand_targets[:, :1], dtype=bool),
                 cand_targets[:, 1:] != cand_targets[:, :-1]], axis=1)
            cell_w = cell_mask * (delta_bg + (1.0 - delta_bg)
                                  * changed.astype(cand_logits.dtype))
        else:
            cell_w = cell_mask
        cell_denom = jnp.maximum(cell_w.sum(axis=2), 1e-6)          # (bs, K)
        bce_per_cell = bce.mean(axis=3)                             # (bs, K, 81), mean over 9 digits
        bce_per_slot = ((bce_per_cell * cell_w).sum(axis=2)
                        / cell_denom)                               # (bs, K), mean over empty cells
        if backtrack:
            # Strict stage-replay: supervise ONLY the stage-t readout = the last
            # active latent slot per example (index k-1), whose candidate target
            # is that stage's grid. Earlier slots and the LM CE are dropped, so
            # this step purely re-derives "apply f exactly t times -> stage-t
            # readout." Injection above still spans all active slots, so the
            # recurrence reaching slot k-1 is intact.
            k_per = active_full.sum(axis=1)                         # (bs,)
            last_idx = (k_per - 1)[:, None]                         # (bs, 1)
            slot_sel = ((jnp.arange(num_slots)[None, :] == last_idx)
                        & active_full)                             # (bs, K)
            slot_active = slot_sel.astype(bce_per_slot.dtype)
        else:
            slot_active = active_full.astype(bce_per_slot.dtype)    # (bs, K)
        aux_denom = jnp.maximum(slot_active.sum(), 1.0)
        avg_aux_loss = (bce_per_slot * slot_active).sum() / aux_denom

        # Backtrack (replay) steps train only from the readout (no LM CE).
        ce_term = 0.0 if backtrack else avg_ce_loss
        leak_term = 0.0 if backtrack else leak_weight * leak_loss
        total_loss = ce_term + aux_weight * avg_aux_loss + leak_term

        return total_loss, (pred_logits, avg_ce_loss, avg_aux_loss, leak_loss,
                            sigma_mean)

    # Compute the learning rate and perform gradient descent
    step = state.step
    lr = learning_rate_fn(step)
    (loss, aux), grads = jax.value_and_grad(loss_fn,
                                            has_aux=True)(state.params)
    pred_logits, ce_loss, aux_loss, leak_loss, sigma_mean = aux
    grads = jax.lax.pmean(grads, "batch")
    metrics = {
        "step": step, "loss": loss, "learning_rate": lr,
        "ce_loss": ce_loss,
        "aux_loss": aux_loss,
        "leak_loss": leak_loss,
        # Mean of the learned per-slot sigma. This is the number that tells us
        # whether the noise survives training or is driven to zero.
        "latent_sigma": sigma_mean,
        "pred_logits": pred_logits, "weights": inputs.shape[0]
    }
    if apply_update:
        new_state = state.apply_gradients(grads=grads)
        return new_state, metrics, pred_logits
    return grads, metrics, pred_logits