code/wavecurriculum_run/train/evaluater.py

"""Evaluation related functions."""

import os

from flax.training import common_utils
import jax
from jax import numpy as jnp
import numpy as np

from train import model

import pdb


def _verbose_eval():
    return os.environ.get("SUDOKU_VERBOSE_EVAL", "0") == "1"


def _lcs_len(a, b):
    """Length of the longest common subsequence of two cell sequences.

    Measures order agreement while tolerating insertions, so a single early
    deviation costs one step instead of desyncing every later comparison the
    way a positional match does.
    """
    prev = [0] * (len(b) + 1)
    for x in a:
        cur = [0]
        for k, y in enumerate(b):
            cur.append(prev[k] + 1 if x == y else max(cur[k], prev[k + 1]))
        prev = cur
    return prev[-1]


def _inst_ce_stats(logp_digits, bits, qcounts):
    """CE(q || p) against the ceiling log|S|, per cell.

    q is the empirical per-cell digit distribution of this puzzle's stored
    instances at this stage: the POST-constraint distribution the data actually
    teaches. That is the whole point of using it instead of Uniform(S). By
    Gibbs, CE(Uniform(S) || p) >= log|S| for every p, so comparing that to
    log|S| is a tautology (it is exactly `excess >= 0`). CE(q || p) instead has
    floor H(q), and H(q) < log|S| strictly whenever the constraints prune
    anything, so the comparison

        CE(q || p)  <=  log|S|          equivalently   KL(q || p) <= log|S| - H(q)

    is both satisfiable and meaningful: the model's divergence from the truth
    must be smaller than the information the constraints carry. It is also a
    complete test on its own, unlike mass/spread, because it punishes every
    failure mode at once -- dropping a candidate that instances do use sends
    -log p(d) up, and leaking mass outside S lowers p on the digits that are
    actually targets.

    A cell is only counted when the bar is above the floor, i.e. when
    H(q) < log|S| strictly. Two kinds of cell are dropped:

      |S| = 1     the ceiling is log 1 = 0 and no finite CE can meet it.
      H(q) = log|S|  q is already uniform on S, so the bar EQUALS the floor and
                  by Gibbs only p == q exactly passes. Asking for that is asking
                  for zero error, not for learning.

    On the real pool this second case is 8-27% of cells depending on stage, so
    leaving them in would cap ceq_ok well below 1 and make any fixed threshold
    mean something different at every stage. The exclusion depends only on q and
    S, never on p, so it cannot be gamed by the model.

    Args:
        logp_digits: (n, 9) log-probabilities of digits 1..9 under the model.
        bits: (n,) int bitmask; bit d-1 set iff digit d is a candidate.
        qcounts: (n, 9) instance-digit counts; column d-1 counts digit d.

    Returns:
        Sums over the countable cells plus the count they divide by.
    """
    out = {"ceq": 0.0, "hq": 0.0, "ceq_ok": 0, "ceq_gap": 0.0, "ceq_count": 0,
           "ceq_all": 0, "ceq_groups": 0, "ceq_skip": 0}
    in_s = ((bits[:, None] >> np.arange(9)) & 1).astype(np.float64)
    size = in_s.sum(axis=1)
    tot = qcounts.sum(axis=1)
    cand = (size >= 2) & (tot > 0)
    if not cand.any():
        return out
    q_all = qcounts[cand].astype(np.float64) / tot[cand][:, None]
    h_all = -(q_all * np.log(np.maximum(q_all, 1e-12))).sum(axis=1)
    l_all = np.log(size[cand])
    # Keep only cells where the constraints left headroom under the ceiling.
    room = (l_all - h_all) > 1e-9
    out["ceq_skip"] = int((~room).sum())
    if not room.any():
        return out
    q = q_all[room]
    lp = logp_digits[cand][room]
    ce_q = -(q * lp).sum(axis=1)
    h_q = h_all[room]
    gap = l_all[room] - ce_q
    out["ceq"] = float(ce_q.sum())
    out["hq"] = float(h_q.sum())
    out["ceq_ok"] = int((gap >= 0).sum())
    out["ceq_gap"] = float(gap.sum())
    out["ceq_count"] = int(room.sum())
    out["ceq_all"] = int(bool((gap >= 0).all()))
    out["ceq_groups"] = 1
    return out


def _cand_value_stats(logp_digits, bits):
    """Value-head statistics against one stage's candidate sets.

    S is the RAW wave-solver candidate set at that cell (the bitmask in
    cand_targets), NOT the support of the filtered instances after constraint
    solving. A digit can be in S and still never appear in a surviving
    instance; we still want p uniform over S, because that is the
    pre-constraint superposition.

    The instance target at a cell is drawn uniformly from S, so the CE the
    trainer minimizes is
        CE = -(1/|S|) sum_{d in S} log p(d)  >=  log|S|,
    with equality iff p == Uniform(S); collapsing onto one candidate sends it
    to infinity. So CE - log|S| == KL(Uniform(S) || p) is zero exactly at
    true superposition and grows from leakage outside S or collapse inside it.

    Args:
        logp_digits: (n, 9) log-probabilities of digits 1..9 under the model.
        bits: (n,) int bitmask; bit d-1 set iff digit d is a candidate.

    Returns:
        Dict of sums over the n cells plus the counts they divide by. `spread`
        (normalized within-set entropy, 1.0 = uniform over S) is only defined
        for |S| >= 2, hence its own count.
    """
    in_s = ((bits[:, None] >> np.arange(9)) & 1).astype(np.float64)  # (n, 9)
    size = in_s.sum(axis=1)
    p_d = np.exp(logp_digits)
    mass = (in_s * p_d).sum(axis=1)
    ce = -(in_s * logp_digits).sum(axis=1) / size
    # excess splits exactly into the two things we want driven to zero:
    #   excess = CE - log|S| = log(1/mass) + KL(Uniform(S) || p/mass)
    #            \_________/   \_________/
    #             leak outside   non-uniformity inside
    # Both terms are >= 0, so excess == 0 iff all mass is on S AND spread
    # uniformly over it. Derivation: writing p(d) = mass * q(d) for d in S,
    # CE = -log(mass) - (1/|S|) sum log q(d) = -log(mass) + log|S| + KL(U||q).
    excess = ce - np.log(size)
    # Entropy of the model's digit distribution, renormalized over digits 1..9
    # (the log_softmax ran over the full vocab, which also holds the latent and
    # coordinate tokens). This is deliberately NOT the `spread` numerator below:
    # that one renormalizes onto S, so it satisfies H <= log|S| by construction
    # and cannot detect leakage. Over all 9 digits H can reach log 9, and
    #     H(p) > log|S|  ==>  p carries mass outside S,
    # since Uniform(S) is the max-entropy distribution supported on S. So the
    # bound is a necessary condition for having learned the stage's constraints,
    # and it tightens as the curriculum shrinks |S|.
    p_dig = p_d / np.maximum(p_d.sum(axis=1, keepdims=True), 1e-12)
    h_full = -(p_dig * np.log(np.maximum(p_dig, 1e-12))).sum(axis=1)
    out = {
        "ce": float(ce.sum()),
        "floor": float(np.log(size).sum()),
        "mass": float(mass.sum()),
        "count": int(size.size),
        "spread": 0.0,
        "spread_count": 0,
        # Same quantities restricted to genuinely superposed cells (|S| >= 2).
        # Singleton cells are already determined at this stage, so there is
        # nothing to be uniform about; as the curriculum deepens they come to
        # dominate the pooled average and dilute the superposition signal.
        "excess_multi": 0.0,
        "mass_multi": 0.0,
        "kl_multi": 0.0,
        # Entropy-ceiling aggregates. The inequality is per cell, and averaging
        # H and log|S| separately is NOT equivalent to it: a mean gap can stay
        # positive while many individual cells violate. So count cells.
        "hbound_ok": 0,       # cells satisfying H(p) <= log|S|
        "hgap": 0.0,          # sum of log|S| - H(p), the margin in nats
        "hbound_all": 0,      # 1 iff every |S|>=2 cell here satisfies it
        "hbound_groups": 0,   # number of (puzzle, stage) groups counted
    }
    multi = size >= 2
    if multi.any():
        m = np.maximum(mass[multi], 1e-12)
        q = in_s[multi] * p_d[multi] / m[:, None]
        h = -(q * np.log(np.maximum(q, 1e-12))).sum(axis=1)
        out["spread"] = float((h / np.log(size[multi])).sum())
        out["spread_count"] = int(multi.sum())
        out["excess_multi"] = float(excess[multi].sum())
        out["mass_multi"] = float(mass[multi].sum())
        # KL(Uniform(S) || q): the non-uniformity term of the split above.
        # Stricter than 1 - spread, because it weights every candidate equally
        # and so blows up when the model drops one candidate to ~0.
        out["kl_multi"] = float((excess[multi] + np.log(m)).sum())
        gap = np.log(size[multi]) - h_full[multi]
        out["hbound_ok"] = int((gap >= 0).sum())
        out["hgap"] = float(gap.sum())
        out["hbound_all"] = int(bool((gap >= 0).all()))
        out["hbound_groups"] = 1
    return out


def _wave_order_score(emitted, truth, waves, max_wave):
    """How many emitted cells came from the earliest still-unfilled wave.

    A cell's wave is the propagation depth that determines it. Cells sharing a
    wave are order-interchangeable, so the solver order in the data is only one
    valid linearization; crediting any cell from the frontier wave measures
    "is the model respecting propagation depth" without penalizing an arbitrary
    tie-break. Cells that repeat or are not real empty cells score nothing.
    """
    left = dict(zip(truth, waves))
    per_wave = np.bincount(np.asarray(waves, dtype=np.int64),
                          minlength=max_wave + 1)
    min_w, ok = 0, 0
    for cell in emitted:
        while min_w <= max_wave and per_wave[min_w] == 0:
            min_w += 1
        if min_w > max_wave:
            break
        cw = left.pop(cell, None)
        if cw is None:
            continue
        per_wave[cw] -= 1
        ok += int(cw == min_w)
    return ok


def valid_solution(output_seq):
    """
    This function checks if the puzzle is a valid solution by verifying if
    each row, column and box has all the numbers from 1 to 9.

    Args:
        output_seq: a numpy array of shape (243,) containing the sequence of
            output numbers

    Returns:
        int: 1 if correct solution, otherwise returns 0
    """
    # rows[i, j] keeps track if ith row has received (j + 1) number
    rows = np.zeros((9, 9))
    # cols[i, j] keeps track if ith column has received (j + 1) number
    cols = np.zeros((9, 9))
    # boxes[i, j] keeps track if ith box has received (j + 1) number
    boxes = np.zeros((9, 9))

    for j in range(81):
        # The row and column are in the range (0, 8) and puzzle values are in (1, 9)
        if int(output_seq[3 * j]) >= 9:
            return False
        if int(output_seq[3 * j + 1]) >= 9:
            return False
        if int(output_seq[3 * j + 2]) > 9:
            return False
        
        row_num = int(output_seq[3 * j])
        col_num = int(output_seq[3 * j + 1])
        
        # Mark the number in the row, column and box
        rows[row_num, int(output_seq[3 * j + 2] - 1)] += 1
        cols[col_num, int(output_seq[3 * j + 2] - 1)] += 1
        boxes[
            int(3 * (row_num // 3) + (col_num // 3)), int(output_seq[3 * j + 2] - 1)
        ] += 1

    if np.all(rows) and np.all(cols) and np.all(boxes):
        return True
    else:
        return False


def eval_step(state, batch, latent_vals, slot_pos, latent_active, config):
    pred_logits, hidden, cand_logits = model.TransformerLMHeadModel(config).apply(
        {"params": state.params}, batch, latent_values=latent_vals,
        latent_positions=slot_pos, latent_active=latent_active,
        )
    return pred_logits, hidden, cand_logits


def verify_sudoku_board(puzzle, row_num, col_num, num):
    """
    Args:
        puzzle (np.array): The correct Sudoku puzzle.
        row_num (int): The row number (0-8).
        col_num (int): The column number (0-8).
        num (int): The number predicted at the specified row and column.

    Raises:
        AssertionError: If the row_num * 9 + col_num >= 81 or if the number at the specified row and column is not equal to the given number.
    """
    if row_num * 9 + col_num >= 81: 
        assert False
  
    assert puzzle[row_num * 9 + col_num] == num


def get_eval_metrics(state, eval_data_iter, p_eval_step, config):
    """This function computes given evaluation metrics (e.g, accuracy) in eval metrics for each batch and appends the metric in the list of eval_metrics.

    Args: 
        state: contains model parameters, optimizer, etc.
        eval_data_iter: data iterator for evaluation dataset
        p_eval_step: pmap function for forward pass of model for evaluation
        config: general experiment config file

    Returns: 
        eval_metrics: contains list of evaluation metrics for each batch
    """

    eval_metrics = {
        "acc": [],  # Unique-solution placement (parent leftover; not printed as val_acc)
        "loc_acc": [],  # Location acc: model picks the ground-truth next cell (r,c)
        # Permutation-tolerant location diagnostics. The upstream code scored
        # only "acc" -- the digit at the cell the model chose -- and never
        # compared (r,c) to the target order, because the solver order is one
        # arbitrary linearization of a partial order: cells that become
        # determined in the same propagation wave are interchangeable. loc_acc
        # above therefore reads ~0.04 even when the model emits a valid
        # permutation of the right cells, since one early deviation desyncs the
        # rest of the positional comparison. These three replace it.
        "loc_coverage": [],  # distinct emitted cells that are really empty / #empty
        "loc_dup": [],       # fraction of emitted cells that repeat an earlier one
        "loc_lcs": [],       # longest common subsequence with target order / #empty
        "loc_wave": [],      # emitted cell sits in the earliest unfilled wave
        "val_given_loc_acc": [],  # Correct digit AMONG steps where location matched
        "cand_bit_acc": [],  # Per-digit accuracy of predicted candidate masks
        "cand_set_acc": [],  # Exact candidate-SET match per empty cell (all 9 bits)
        "cand_set_acc_changed": [],  # ...restricted to cells that changed this stage
        "acc_complete_puzzle": []  # Accuracy of predicting correct complete puzzle
    }
    # Per-difficulty-level cell accuracy (levels 3..8). Diagnostic only: the
    # curriculum no longer keys on level.
    level_ok = {lvl: 0 for lvl in range(3, 9)}
    level_tot = {lvl: 0 for lvl in range(3, 9)}

    K = int(config.num_latent_slots)

    # Per-SLOT candidate-set accuracy, i.e. per reasoning depth. Slot j holds
    # wave snapshot j, so slot_ok[j]/slot_tot[j] is "how well is propagation
    # block j predicted". This is the signal the depth curriculum promotes and
    # backtracks on, replacing the old per-level accuracy.
    slot_ok = np.zeros(max(K, 1), dtype=np.int64)
    slot_tot = np.zeros(max(K, 1), dtype=np.int64)
    slot_ok_ch = np.zeros(max(K, 1), dtype=np.int64)
    slot_tot_ch = np.zeros(max(K, 1), dtype=np.int64)

    # Per-stage in-set rate: was the emitted digit a *member* of that stage's
    # candidate set? This is the promotion signal for the instance arm, where
    # the target is one sampled assignment rather than the unique solution, so
    # the model is right to emit any candidate. The candidate masks are read
    # here as a metric only; nothing supervises them.
    inset_ok = np.zeros(max(K, 1), dtype=np.int64)
    inset_tot = np.zeros(max(K, 1), dtype=np.int64)

    # Per-stage value-distribution stats, measured TEACHER-FORCED so the cell is
    # given and the value head is scored in isolation from the model's choice of
    # location. For a cell whose stage-s candidate set is S, the instance target
    # is drawn uniformly from S, so the CE the trainer minimizes is
    #     CE = -(1/|S|) sum_{d in S} log p(d)  >=  log|S|,
    # with equality iff p == Uniform(S). Collapsing onto a single candidate
    # sends CE to infinity. So (CE - log|S|) == KL(Uniform(S) || p) is a single
    # number that is 0 exactly at true superposition and grows from either
    # leakage outside S or mode collapse inside it.
    vce_sum = np.zeros(max(K, 1), dtype=np.float64)    # sum of CE vs Uniform(S)
    vfloor_sum = np.zeros(max(K, 1), dtype=np.float64)  # sum of log|S|
    vmass_sum = np.zeros(max(K, 1), dtype=np.float64)   # sum of prob mass on S
    vspread_sum = np.zeros(max(K, 1), dtype=np.float64)  # sum of H(p|S)/log|S|
    vcnt = np.zeros(max(K, 1), dtype=np.int64)
    vspread_cnt = np.zeros(max(K, 1), dtype=np.int64)   # only |S| >= 2 cells
    # Same, restricted to |S| >= 2 (the cells that carry real superposition).
    vexc_m_sum = np.zeros(max(K, 1), dtype=np.float64)
    vmass_m_sum = np.zeros(max(K, 1), dtype=np.float64)
    vkl_m_sum = np.zeros(max(K, 1), dtype=np.float64)
    # Entropy ceiling H(p) <= log|S|, counted per cell and per (puzzle, stage).
    vhok_sum = np.zeros(max(K, 1), dtype=np.int64)
    vhgap_sum = np.zeros(max(K, 1), dtype=np.float64)
    vhall_sum = np.zeros(max(K, 1), dtype=np.int64)
    vhgrp_sum = np.zeros(max(K, 1), dtype=np.int64)
    # CE(q||p) vs the ceiling log|S|, q = the instances' own digit frequencies.
    vceq_sum = np.zeros(max(K, 1), dtype=np.float64)
    vhq_sum = np.zeros(max(K, 1), dtype=np.float64)
    vceq_ok = np.zeros(max(K, 1), dtype=np.int64)
    vceq_gap = np.zeros(max(K, 1), dtype=np.float64)
    vceq_cnt = np.zeros(max(K, 1), dtype=np.int64)
    vceq_all = np.zeros(max(K, 1), dtype=np.int64)
    vceq_grp = np.zeros(max(K, 1), dtype=np.int64)

    # Per-round-bin cell accuracy, for the round-count DATA curriculum: bin b is
    # unlocked at stage b, so bin_ok[b]/bin_tot[b] measures competence on the
    # puzzles that stage b introduced. This is the promotion signal for the arm
    # that has no latent slots and therefore no per-depth signal.
    n_bins = int(getattr(config, "curriculum_max_stage", 12))
    bin_ok = {b: 0 for b in range(1, n_bins + 1)}
    bin_tot = {b: 0 for b in range(1, n_bins + 1)}

    for eval_epoch in range(config.eval_epochs):
        with jax.profiler.StepTraceAnnotation("eval", step_num=eval_epoch):

            batch_tuple = next(eval_data_iter)

            # Input seq is (batchsize, 3*81 + K): clue triples, K latent
            # placeholder slots, then solution triples.
            input_seq = np.array(batch_tuple[0])

            # Puzzle solution is of the shape (batchsize, 81). Each pos in {0,.., 80}
            # for each puzzle contains value at cell (pos//9+1, pos%9 + 1)
            puzzle_sol = np.array(batch_tuple[1])
            start_index = np.array(batch_tuple[2])
            levels = np.array(batch_tuple[3]).reshape(-1)
            rbins = (np.array(batch_tuple[5]).reshape(-1)
                     if len(batch_tuple) > 5 else np.zeros_like(levels))
            # (bs, K, 81, 9) instance-digit counts; zeros when the run has no
            # instance pool, in which case CE(q||p) is simply not reported.
            q_counts = (np.array(batch_tuple[6]).astype(np.int64)
                        if len(batch_tuple) > 6 else None)
            total_pred, sucess_pred = 0, 0
            # Location = did the model emit the ground-truth next (r,c) cell.
            loc_tot, loc_ok, val_given_loc_ok = 0, 0, 0

            bs = input_seq.shape[0]
            bidx = np.arange(bs)
            si3 = 3 * start_index.reshape(-1)
            slot_pos = si3[:, None] + np.arange(K)[None, :]
            if getattr(config, "cand_slot_mode", "level") == "depth":
                # Eval always builds all K latents, so score all K slots.
                k_budget = np.full_like(levels, K)
            else:
                k_budget = np.clip(levels - 2, 1, K)
            active_full = np.arange(K)[None, :] < k_budget[:, None]

            def run_model(seq_batch, latent_vals, act, want_cand=False):
                sharded = common_utils.shard(
                    jax.tree_util.tree_map(np.asarray, seq_batch))
                # Explicit reshape so a zero-width slot dim (K=0 baseline)
                # shards without the ambiguous -1 inference of shard().
                _nd = jax.local_device_count()
                def _shard(x):
                    x = np.asarray(x)
                    return x.reshape((_nd, x.shape[0] // _nd) + x.shape[1:])
                lv = _shard(latent_vals)
                lp = _shard(slot_pos)
                la = _shard(act)
                logits, hidden, cand = p_eval_step(state, sharded, lv, lp, la)
                logits = np.array(logits).reshape(bs, *np.array(logits).shape[2:])
                hidden = np.array(hidden).reshape(bs, *np.array(hidden).shape[2:])
                if want_cand:
                    cand = np.array(cand).reshape(bs, *np.array(cand).shape[2:])
                    return logits, hidden, cand
                return logits, hidden

            # ---- Build the continuous latent thoughts (K recurrence passes,
            # difficulty-matched budget; causal masking means only the clue
            # region influences them). ----
            latent_vals = np.zeros((bs, K, config.emb_dim), dtype=np.float32)
            build_seq = np.array(input_seq)
            build_seq_masked = np.array(build_seq)
            # Hide the solution region during latent build (safety; causality
            # already prevents leakage into slot hiddens).
            for j in range(bs):
                build_seq_masked[j, si3[j] + K:] = 0
            # Recurrent feedback: build each latent thought from the previous
            # slot's hidden. Skipped when the model does not inject latents
            # (no-recurrence control): slots stay as static placeholders, so
            # latent_vals is left at zeros and never used.
            recurrent = bool(int(getattr(config, "recurrent_latent", 1)))
            if recurrent and K > 0:
                for j in range(K):
                    act_j = active_full & (np.arange(K)[None, :] < j)
                    _, hidden = run_model(build_seq_masked, latent_vals, act_j)
                    src = si3 - 1 + j
                    latent_vals[:, j] = hidden[bidx, src]

            # ---- Candidate-set prediction accuracy (the multi-value target) ----
            # One forward pass with the fully-built latents; read the per-slot
            # candidate head and compare to the staged bitmask targets, scored
            # only over active slots and empty cells (clue cells were zeroed).
            # Skipped entirely for the K=0 no-latent baseline (no candidate head).
            pred_bits = tgt_bits = cand_targets = None
            if K > 0:
                cand_targets = np.array(batch_tuple[4]).astype(np.int64)   # (bs, K, 81)
            # The candidate head is off in the instance arm (aux weight 0), so
            # skip its forward pass and set metrics; the masks above are still
            # read for the in-set rate.
            if K > 0 and float(getattr(config, "aux_cand_weight", 1.0)) > 0.0:
                _, _, cand_logits = run_model(
                    build_seq_masked, latent_vals, active_full, want_cand=True)  # (bs,K,81,9)
                pred_bits = (np.array(cand_logits) > 0.0)                  # sigmoid>0.5
                tgt_bits = ((cand_targets[..., None] >> np.arange(9)) & 1).astype(bool)
                valid = (cand_targets > 0) & active_full[:, :, None]       # (bs,K,81)
                if valid.sum() > 0:
                    bit_match = (pred_bits == tgt_bits)                    # (bs,K,81,9)
                    eval_metrics["cand_bit_acc"].append(
                        float(bit_match[valid].mean()))
                    eval_metrics["cand_set_acc"].append(
                        float(bit_match.all(axis=3)[valid].mean()))
                    # Same score restricted to cells whose candidate set
                    # actually changed from the previous stage. The unrestricted
                    # metrics above are dominated by cells that are unchanged
                    # copies of slot j-1, so they stay high for a head that has
                    # learned nothing but "repeat the previous slot".
                    changed = np.concatenate(
                        [np.ones_like(cand_targets[:, :1], dtype=bool),
                         cand_targets[:, 1:] != cand_targets[:, :-1]], axis=1)
                    valid_ch = valid & changed
                    if valid_ch.sum() > 0:
                        eval_metrics["cand_set_acc_changed"].append(
                            float(bit_match.all(axis=3)[valid_ch].mean()))
                    # Accumulate the same score split by slot (= depth).
                    set_match = bit_match.all(axis=3)                   # (bs,K,81)
                    slot_ok += (set_match & valid).sum(axis=(0, 2))
                    slot_tot += valid.sum(axis=(0, 2))
                    slot_ok_ch += (set_match & valid_ch).sum(axis=(0, 2))
                    slot_tot_ch += valid_ch.sum(axis=(0, 2))

            # ---- Teacher-forced value distribution vs the candidate set ----
            # One forward pass on the ground-truth solver-order sequence. The
            # prefix pins down which cell each value slot refers to, so this
            # measures the value head alone: location cannot contaminate it,
            # and no sampling is needed because the full softmax is available.
            if K > 0 and cand_targets is not None:
                tf_logits, _ = run_model(input_seq, latent_vals, active_full)
                # log_softmax over the whole vocab, matching the training CE.
                tf_shift = tf_logits - tf_logits.max(axis=-1, keepdims=True)
                tf_logp = tf_shift - np.log(
                    np.exp(tf_shift).sum(axis=-1, keepdims=True))
                for j in range(bs):
                    si = int(start_index[j].reshape(-1)[0])
                    n_out = 81 - si
                    if n_out <= 0:
                        continue
                    base = 3 * si + K
                    t = np.arange(n_out)
                    v_pos = base + 3 * t + 2          # value token positions
                    if v_pos[-1] >= config.seq_len:
                        keep = v_pos < config.seq_len
                        t, v_pos = t[keep], v_pos[keep]
                        if t.size == 0:
                            continue
                    # Token ids are the numbers themselves: rows/cols 0..8 and
                    # digits 1..9, so digit d lives at vocab index d.
                    cells = (input_seq[j, base + 3 * t] * 9
                             + input_seq[j, base + 3 * t + 1]).astype(np.int64)
                    ok_cell = (cells >= 0) & (cells < 81)
                    # logits at index p-1 predict the token at position p.
                    logp_d = tf_logp[j, v_pos - 1, 1:10]        # (n, 9)
                    for s in range(K):
                        bits = cand_targets[j, s, np.where(ok_cell, cells, 0)]
                        # bits <= 0 marks a clue cell, which is not supervised.
                        sel = ok_cell & (bits > 0)
                        if not sel.any():
                            continue
                        st = _cand_value_stats(logp_d[sel], bits[sel])
                        vce_sum[s] += st["ce"]
                        vfloor_sum[s] += st["floor"]
                        vmass_sum[s] += st["mass"]
                        vcnt[s] += st["count"]
                        vspread_sum[s] += st["spread"]
                        vspread_cnt[s] += st["spread_count"]
                        vexc_m_sum[s] += st["excess_multi"]
                        vmass_m_sum[s] += st["mass_multi"]
                        vkl_m_sum[s] += st["kl_multi"]
                        vhok_sum[s] += st["hbound_ok"]
                        vhgap_sum[s] += st["hgap"]
                        vhall_sum[s] += st["hbound_all"]
                        vhgrp_sum[s] += st["hbound_groups"]
                        if q_counts is not None:
                            qs = _inst_ce_stats(
                                logp_d[sel], bits[sel],
                                q_counts[j, s, np.where(ok_cell, cells, 0)][sel])
                            vceq_sum[s] += qs["ceq"]
                            vhq_sum[s] += qs["hq"]
                            vceq_ok[s] += qs["ceq_ok"]
                            vceq_gap[s] += qs["ceq_gap"]
                            vceq_cnt[s] += qs["ceq_count"]
                            vceq_all[s] += qs["ceq_all"]
                            vceq_grp[s] += qs["ceq_groups"]

            min_start_index = int(np.min(start_index))
            cur_input_seq = input_seq[:, :(min_start_index*3)]
            for i in range(min_start_index * 3, config.seq_len):
                ### In i^th iteration, i^th number in sequence will predict
                padding = np.zeros((input_seq.shape[0],
                                    config.seq_len - len(cur_input_seq[0])),
                                dtype=np.int32)
                concat_batch = np.hstack((cur_input_seq, padding))

                pred_logits, _ = run_model(concat_batch, latent_vals, active_full)

                # Positions < 3*start_index + K are given (clues + latent
                # slots); the model predicts from there on. K is a multiple
                # of 3, so the triple phase of i is unchanged.
                if i%3 == 2:
                    # Model predicts the value at the cell (cur_input_seq[j][i-2],
                    # cur_input_seq[j][i-1])
                    max_number = pred_logits[:, i-1, :].argmax(axis=-1).flatten()
                    mask_arr = np.array(i >= (3 * start_index + K)).squeeze()

                    next_number = max_number * mask_arr + (1 - mask_arr) * input_seq[:, i]

                    cur_input_seq = np.hstack(
                        (cur_input_seq, np.reshape(next_number, (-1, 1)))
                    )

                    # Iterate through all examples in batch and calculate successful
                    # predictions of numbers
                    for j in range(len(cur_input_seq)):
                        if not mask_arr[j]:
                            continue

                        total_pred += 1
                        level_tot[int(levels[j])] += 1
                        if int(rbins[j]) in bin_tot:
                            bin_tot[int(rbins[j])] += 1

                        # Location accuracy: did the model emit the ground-truth
                        # next cell (r,c) for this solver-order step?
                        loc_tot += 1
                        loc_match = (int(cur_input_seq[j][i-2]) == int(input_seq[j, i-2])
                                     and int(cur_input_seq[j][i-1]) == int(input_seq[j, i-1]))
                        if loc_match:
                            loc_ok += 1

                        # In-set rate per stage, scored at the ground-truth cell
                        # so a wrong location cannot make a digit vacuously
                        # legal. cand_targets[j, s, cell] is stage s's bitmask
                        # under cand_slot_mode="depth" (slot s <-> stage s).
                        if cand_targets is not None and loc_match:
                            cell = (int(input_seq[j, i-2]) * 9
                                    + int(input_seq[j, i-1]))
                            v = int(cur_input_seq[j][i])
                            for s in range(K):
                                bits = int(cand_targets[j, s, cell])
                                if bits <= 0:      # clue cell, not supervised
                                    continue
                                inset_tot[s] += 1
                                if 1 <= v <= 9 and (bits >> (v - 1)) & 1:
                                    inset_ok[s] += 1

                        try:
                            verify_sudoku_board(puzzle_sol[j], cur_input_seq[j][i-2], 
                                                cur_input_seq[j][i-1], cur_input_seq[j][i])
                        except AssertionError:
                            # Mistake
                            pass
                        else:
                            sucess_pred += 1
                            level_ok[int(levels[j])] += 1
                            if int(rbins[j]) in bin_ok:
                                bin_ok[int(rbins[j])] += 1
                            if loc_match:
                                val_given_loc_ok += 1
                else:
                    # Model predicts either a row number or column number
                    max_pos = pred_logits[:, i-1, :].argmax(axis=-1).flatten()
                    mask = (i >= (3 * start_index + K)).squeeze()
                    next_pos = max_pos * mask + (1 - mask) * input_seq[:, i]
                    
                    # pdb.set_trace()
                    cur_input_seq = np.hstack(
                        (cur_input_seq, np.reshape(next_pos, (-1, 1)))
                    )

            eval_metrics["acc"].append(sucess_pred * 1.0/ total_pred)
            eval_metrics["loc_acc"].append(loc_ok * 1.0 / max(loc_tot, 1))
            eval_metrics["val_given_loc_acc"].append(
                val_given_loc_ok * 1.0 / max(loc_ok, 1))

            def strip_latent_slots(seq, si):
                return np.concatenate([seq[:3*si], seq[3*si + K:]])

            # ---- Permutation-tolerant location diagnostics ----
            # A cell's wave = the first stage at which its candidate set becomes
            # a singleton, i.e. the propagation depth that determines it. Cells
            # sharing a wave are order-interchangeable, so "did you name a cell
            # from the earliest wave that is still unfilled" is the meaningful
            # ordering signal; the exact index within the wave is arbitrary.
            cov_b, dup_b, lcs_b, wav_b = [], [], [], []
            for j in range(bs):
                si = int(start_index[j].reshape(-1)[0])
                pred = strip_latent_slots(cur_input_seq[j], si)
                true = strip_latent_slots(input_seq[j], si)
                emitted = [(int(pred[3 * k]), int(pred[3 * k + 1]))
                           for k in range(si, 81)]
                truth = [(int(true[3 * k]), int(true[3 * k + 1]))
                         for k in range(si, 81)]
                if not truth or not emitted:
                    continue
                n_true = len(truth)
                cov_b.append(len(set(emitted) & set(truth)) / n_true)
                dup_b.append(1.0 - len(set(emitted)) / len(emitted))
                # LCS is quadratic, so sample a few examples per batch.
                if j < 32:
                    lcs_b.append(_lcs_len(emitted, truth) / n_true)

                if cand_targets is None or K == 0:
                    continue
                cids = np.array([r * 9 + c for r, c in truth], dtype=np.int64)
                bits = cand_targets[j][:, cids]                     # (K, n)
                singleton = (bits > 0) & ((bits & (bits - 1)) == 0)
                w = np.where(singleton.any(axis=0), singleton.argmax(axis=0), K)
                wav_b.append(
                    _wave_order_score(emitted, truth, w, K) / n_true)

            if cov_b:
                eval_metrics["loc_coverage"].append(float(np.mean(cov_b)))
                eval_metrics["loc_dup"].append(float(np.mean(dup_b)))
            if lcs_b:
                eval_metrics["loc_lcs"].append(float(np.mean(lcs_b)))
            if wav_b:
                eval_metrics["loc_wave"].append(float(np.mean(wav_b)))

            # ---- Print one concrete example answer the model generated ----
            if eval_epoch == 0:
                j = 0
                si = int(start_index[j, 0])
                pred = strip_latent_slots(cur_input_seq[j], si)
                shown, n_ok, n_tot = [], 0, 0
                for k in range(si, 81):
                    r, c, v = int(pred[3*k]), int(pred[3*k+1]), int(pred[3*k+2])
                    tv = int(puzzle_sol[j][r*9+c]) if (0 <= r < 9 and 0 <= c < 9) else -1
                    ok = (0 <= r < 9 and 0 <= c < 9 and v == tv)
                    n_tot += 1; n_ok += int(ok)
                    if len(shown) < 12:
                        shown.append(f"({r},{c})->{v}[true {tv}]{'ok' if ok else 'X'}")
                if _verbose_eval():
                    print(f"EXAMPLE (level={int(levels[j])}, k={int(k_budget[j])}): "
                          f"model emitted {n_tot} (r,c)->v triples for the empty cells "
                          f"(format: (row,col)->value[true T]); first 12:", flush=True)
                    print("   ", " ".join(shown), flush=True)
                    print(f"EXAMPLE cells-correct={n_ok}/{n_tot}  "
                          f"valid_full_grid={valid_solution(pred)}", flush=True)

                # Instance arm: emitted digit next to the deepest stage's
                # candidate set, so it is visible whether the model is sitting
                # inside the superposition or outside it.
                if K > 0 and cand_targets is not None and pred_bits is None:
                    tgt = strip_latent_slots(input_seq[j], si)
                    shown = []
                    for t3 in range(si, min(si + 8, 81)):
                        r, c = int(tgt[3*t3]), int(tgt[3*t3+1])
                        bits = int(cand_targets[j, K-1, r*9+c])
                        cset = "".join(str(d+1) for d in range(9)
                                       if (bits >> d) & 1)
                        shown.append(f"(r{r},c{c})->{int(pred[3*t3+2])} "
                                     f"in{{{cset}}}")
                    if _verbose_eval():
                        print(f"EXAMPLE emitted vs stage-{K} candidate set:",
                              " ".join(shown), flush=True)

                # ---- Candidate-set (multi-value) prediction for this puzzle ----
                # Show, at the last active latent slot, predicted vs target
                # candidate SETS for the first few empty cells. (No latent
                # slots in the K=0 baseline, so nothing to show.)
                if K > 0 and pred_bits is not None:
                    kj = int(k_budget[j]) - 1
                    def _digs(bitrow):
                        return "".join(str(d + 1) for d in range(9) if bitrow[d])
                    cand_shown = []
                    for cell in range(81):
                        if cand_targets[j, kj, cell] <= 0:   # clue / not supervised
                            continue
                        r, c = cell // 9, cell % 9
                        pset = _digs(pred_bits[j, kj, cell])
                        tset = _digs(tgt_bits[j, kj, cell])
                        cand_shown.append(f"(r{r},c{c}) pred{{{pset}}} true{{{tset}}}")
                        if len(cand_shown) >= 8:
                            break
                    if _verbose_eval():
                        print(f"EXAMPLE candidate-set @slot{kj} (pred vs true):",
                              " ".join(cand_shown), flush=True)

            correct_eval_sudoku_puzzle = 0

            for i in range(len(cur_input_seq)):

                # increase correct_eval_sudoku_puzzle when the model output solution
                # for a given puzzle is correct
                stripped = strip_latent_slots(cur_input_seq[i], int(start_index[i, 0]))
                correct_eval_sudoku_puzzle += valid_solution(stripped)

            eval_metrics["acc_complete_puzzle"].append(
                correct_eval_sudoku_puzzle * 1.0 / len(cur_input_seq)
            )

    per_level = {lvl: (level_ok[lvl] / level_tot[lvl] if level_tot[lvl] else -1.0)
                 for lvl in range(3, 9)}
    eval_metrics["per_level_acc"] = per_level
    if _verbose_eval():
        print("PER-LEVEL cell acc:",
              {lvl: (f"{v:.3f}" if v >= 0 else "n/a") for lvl, v in per_level.items()},
              flush=True)

    # Per-depth candidate-set accuracy, keyed by stage (slot j -> stage j+1) so
    # the curriculum controller can index it directly by stage number.
    per_slot = {j + 1: (float(slot_ok[j] / slot_tot[j]) if slot_tot[j] else -1.0)
                for j in range(K)}
    per_slot_ch = {j + 1: (float(slot_ok_ch[j] / slot_tot_ch[j])
                           if slot_tot_ch[j] else -1.0) for j in range(K)}
    eval_metrics["per_slot_acc"] = per_slot
    eval_metrics["per_slot_acc_changed"] = per_slot_ch

    # Keyed by stage (slot s -> stage s+1) to match per_slot_acc.
    per_stage_inset = {s + 1: (float(inset_ok[s] / inset_tot[s])
                               if inset_tot[s] else -1.0) for s in range(K)}
    eval_metrics["per_stage_inset_acc"] = per_stage_inset

    # Per-stage value-distribution metrics, keyed by stage to match the above.
    # val_excess = CE(uniform-over-candidates || model) - log|S| >= 0 is the
    # superposition score: 0 means the model spreads exactly uniformly over the
    # stage's candidate set, and it rises if probability leaks outside the set
    # or collapses onto one member of it.
    def _per_stage(num, den):
        return {s + 1: (float(num[s] / den[s]) if den[s] else -1.0)
                for s in range(K)}

    per_stage_vce = _per_stage(vce_sum, vcnt)
    per_stage_vfloor = _per_stage(vfloor_sum, vcnt)
    per_stage_vexcess = {
        s: (per_stage_vce[s] - per_stage_vfloor[s]
            if per_stage_vce[s] >= 0 else -1.0) for s in per_stage_vce}
    eval_metrics["per_stage_val_ce"] = per_stage_vce
    eval_metrics["per_stage_val_floor"] = per_stage_vfloor
    eval_metrics["per_stage_val_excess"] = per_stage_vexcess
    eval_metrics["per_stage_val_mass"] = _per_stage(vmass_sum, vcnt)
    eval_metrics["per_stage_val_spread"] = _per_stage(vspread_sum, vspread_cnt)
    # Superposition metrics on |S| >= 2 cells only. excess_multi is the single
    # number to drive to 0: it equals log(1/mass_multi) + kl_multi, so it falls
    # only when leakage outside the set AND non-uniformity inside it both fall.
    eval_metrics["per_stage_val_excess_multi"] = _per_stage(
        vexc_m_sum, vspread_cnt)
    eval_metrics["per_stage_val_mass_multi"] = _per_stage(
        vmass_m_sum, vspread_cnt)
    eval_metrics["per_stage_val_kl_multi"] = _per_stage(vkl_m_sum, vspread_cnt)
    # Entropy ceiling, three views of the same per-cell inequality:
    #   hbound     fraction of |S|>=2 cells with H(p) <= log|S|   <- gate on this
    #   hgap       mean margin log|S| - H(p) in nats, <0 means violated on average
    #   hbound_puz fraction of (puzzle, stage) groups where EVERY such cell passes
    # The cell fraction is the right gate: the puzzle view compounds (~50 cells,
    # so 95% per cell leaves ~8% of puzzles clean) and the mean margin hides
    # individual violations behind the cells that pass comfortably.
    eval_metrics["per_stage_val_hbound"] = _per_stage(vhok_sum, vspread_cnt)
    eval_metrics["per_stage_val_hgap"] = _per_stage(vhgap_sum, vspread_cnt)
    eval_metrics["per_stage_val_hbound_puz"] = _per_stage(vhall_sum, vhgrp_sum)
    # The single promotion criterion: CE(q||p) <= log|S| per cell, counted.
    #   ceq      mean CE(q||p), to compare against floor H(q) and ceiling log|S|
    #   hq       mean H(q), the best CE any model could achieve on this data
    #   ceq_ok   fraction of |S|>=2 cells clearing the ceiling  <- gate on this
    #   ceq_gap  mean slack log|S| - CE(q||p); negative means worse than uniform
    #   ceq_puz  fraction of puzzles where every such cell clears it
    # This subsumes mass and spread: leak lowers p on real targets, and dropping
    # a candidate the instances use sends -log p(d) up, so both are penalized by
    # the one number. log|S| needs no tuning and recalibrates per stage.
    eval_metrics["per_stage_val_ceq"] = _per_stage(vceq_sum, vceq_cnt)
    eval_metrics["per_stage_val_hq"] = _per_stage(vhq_sum, vceq_cnt)
    eval_metrics["per_stage_val_ceq_ok"] = _per_stage(vceq_ok, vceq_cnt)
    eval_metrics["per_stage_val_ceq_gap"] = _per_stage(vceq_gap, vceq_cnt)
    eval_metrics["per_stage_val_ceq_puz"] = _per_stage(vceq_all, vceq_grp)
    if _verbose_eval() and K > 0 and any(v >= 0 for v in per_stage_vce.values()):
        print("PER-STAGE val_excess (0 = uniform over candidate set):",
              {s: (f"{v:.3f}" if per_stage_vce[s] >= 0 else "n/a")
               for s, v in per_stage_vexcess.items()}, flush=True)

    if _verbose_eval() and K > 0 and any(v >= 0 for v in per_stage_inset.values()):
        print("PER-STAGE in-set rate (emitted digit is a stage-s candidate):",
              {s: (f"{v:.3f}" if v >= 0 else "n/a")
               for s, v in per_stage_inset.items()}, flush=True)

    per_bin = {b: (bin_ok[b] / bin_tot[b] if bin_tot[b] else -1.0)
               for b in range(1, n_bins + 1)}
    eval_metrics["per_bin_acc"] = per_bin
    if _verbose_eval() and any(v >= 0 for v in per_bin.values()):
        print("PER-ROUND-BIN cell acc:",
              {b: (f"{v:.3f}" if v >= 0 else "n/a") for b, v in per_bin.items()},
              flush=True)
    if _verbose_eval() and K > 0:
        print("PER-DEPTH cand-set acc:",
              {s: (f"{v:.3f}" if v >= 0 else "n/a") for s, v in per_slot.items()},
              flush=True)
        print("PER-DEPTH cand-set acc (changed cells only):",
              {s: (f"{v:.3f}" if v >= 0 else "n/a")
               for s, v in per_slot_ch.items()}, flush=True)

    return eval_metrics