"""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] @jax.jit def _latent_mean(mu_w, mu_b, z): """Mean of the per-slot Gaussian, for the stochastic-latent arm. Training draws the injected thought from N(mu, sigma) where mu is a learned residual on the recurrence output. Eval has to read the same object, so it substitutes mu for the raw thought; otherwise every gate would be scoring a latent the model was never trained at. This is the mean readout, not the full noise marginal. """ z = z.astype(jnp.float32) return z + jnp.einsum("bkd,kde->bke", z, mu_w) + mu_b[None] def _latent_mean_slot(mu_w, mu_b, z, slot): """Mean path when creating slot `slot` from its last-layer hidden.""" z = np.asarray(z, dtype=np.float32) w = np.asarray(mu_w[slot]) b = np.asarray(mu_b[slot]) return z + z @ w + b def _nucleus_size(p, tau): """Smallest number of digits holding `tau` of the mass in a 9-vector.""" p = np.asarray(p, dtype=np.float64).reshape(-1) order = np.argsort(-p) csum = np.cumsum(p[order]) k = int(np.searchsorted(csum, tau) + 1) return min(max(k, 1), p.size) def _inst_ce_stats(logp_digits, bits, qcounts, eps=0.0): """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. Two populations are accumulated, because the stage-level question and the per-cell question are not the same test. POPULATION (the *_pop keys, and the promotion gate): every |S| >= 2 cell that has instances. The gate is the stage average mean CE(q||p) <= mean log|S|, which is exactly "the value cross-entropy this stage is minimizing has been driven below the entropy of the uniform candidate set". Nothing is excluded, because an average has no per-cell feasibility problem: a cell with no headroom simply has to be paid for by cells that do have some. HEADROOM SUBSET (ceq_ok / ceq_gap): the same cells minus those where the bar equals the floor, i.e. H(q) = log|S| (q already uniform on S, so by Gibbs only p == q exactly could pass) and |S| = 1 (ceiling is log 1 = 0). That subset is the right denominator for a *per-cell* pass rate, since 8-27% of cells depending on stage would otherwise be unpassable by construction and cap the rate below 1. It is a strictly harder test than the average and is reported as a diagnostic. The exclusion depends only on q and S, never on p. 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 each population plus the counts they divide by. """ out = {"ceq": 0.0, "hq": 0.0, "ceq_ok": 0, "ceq_gap": 0.0, "ceq_count": 0, "puz_ok": 0, "puz_count": 0, "ceq_skip": 0, "ceq_pop": 0.0, "hq_pop": 0.0, "logm_pop": 0.0, "pop_count": 0, "hinge": 0.0, "hinge_floor": 0.0, "klq": 0.0, "hinge_eps": 0.0, "hinge_cnt": 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]) ce_all = -(q_all * logp_digits[cand]).sum(axis=1) # Population sums: the stage-average gate compares these two directly. out["ceq_pop"] = float(ce_all.sum()) out["hq_pop"] = float(h_all.sum()) out["logm_pop"] = float(l_all.sum()) out["pop_count"] = int(cand.sum()) # Per-puzzle hinge, the same comparison as the gate but resolved on THIS # puzzle before any averaging across puzzles. Because max(0,.) is convex, # averaging these over puzzles is >= the violation of the pooled average # (Jensen), so it is a strictly stronger test: a puzzle still above its own # bar can no longer be paid for by puzzles that have slack to spare. out["hinge"] = float(max(0.0, ce_all.mean() - l_all.mean())) # What the hinge would still read at p == q, i.e. for a model that has # matched the instance distribution exactly. Depends only on q and S, never # on p, so it is the floor no model can beat and the only defensible way to # set a tolerance on `hinge`. Puzzles whose constraints barely bite have # H(q) ~ log|S| and contribute a floor > 0 all by themselves. out["hinge_floor"] = float(max(0.0, h_all.mean() - l_all.mean())) # Per-puzzle KL(q||p) = CE(q||p) - H(q), same averaging as `hinge`. The # difference that matters: `hinge` is cleared by any p that spreads over S, # because CE(q||U_S) = log|S| exactly, so it only tests the SUPPORT of q. # This one is zero only at p == q on S, and uniform-on-S pays the full # log|S| - H(q) instead of passing for free, so it is the statistic to gate # on when the target is the instance distribution itself. out["klq"] = float(max(0.0, ce_all.mean() - h_all.mean())) # The hinge against a tightened bar, CE <= log|S| - eps. eps = 0 is the # plain hinge; eps = log|S| - H(q) would demand p == q exactly, so eps is # only meaningful below the headroom, which varies by stage. out["hinge_eps"] = float(max(0.0, ce_all.mean() - l_all.mean() + eps)) out["hinge_cnt"] = 1 # Per-cell pass rate, over the cells where passing is possible at all. room = (l_all - h_all) > 1e-9 out["ceq_skip"] = int((~room).sum()) if not room.any(): return out gap = l_all[room] - ce_all[room] out["ceq"] = float(ce_all[room].sum()) out["hq"] = float(h_all[room].sum()) out["ceq_ok"] = int((gap >= 0).sum()) out["ceq_gap"] = float(gap.sum()) out["ceq_count"] = int(room.sum()) out["puz_ok"] = int(bool((gap >= 0).all())) out["puz_count"] = 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, curriculum=None): """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 curriculum: current CurriculumState, or None. When given, the stage it names gets a second CE(q||p) pass scored the way training feeds the model -- instance-valued prefix, that stage's latent depth -- and reported under the *_inst keys. The promotion gate reads those. 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) # Population sums over every |S|>=2 cell: the stage-average promotion gate. vceq_pop = np.zeros(max(K, 1), dtype=np.float64) vhq_pop = np.zeros(max(K, 1), dtype=np.float64) vlogm_pop = np.zeros(max(K, 1), dtype=np.float64) vpop_cnt = np.zeros(max(K, 1), dtype=np.int64) # Per-puzzle hinge mean(max(0, CE_i - log|S|_i)), and the p==q floor it can # never go below. Denominator is puzzles, not cells. vhinge_sum = np.zeros(max(K, 1), dtype=np.float64) vhinge_floor = np.zeros(max(K, 1), dtype=np.float64) vhinge_cnt = np.zeros(max(K, 1), dtype=np.int64) # Per-puzzle KL(q||p), the same population as the hinge. vklq_sum = np.zeros(max(K, 1), dtype=np.float64) vheps_sum = np.zeros(max(K, 1), dtype=np.float64) promote_eps = float(getattr(config, "promote_eps", 0.0)) # Same statistics from the train-matched pass (instance prefix, stage depth). # Only the gate stage is filled; the rest stay at -1 via _per_stage's count. vceq_pop_i = np.zeros(max(K, 1), dtype=np.float64) vlogm_pop_i = np.zeros(max(K, 1), dtype=np.float64) vhq_pop_i = np.zeros(max(K, 1), dtype=np.float64) vpop_cnt_i = np.zeros(max(K, 1), dtype=np.int64) vhinge_sum_i = np.zeros(max(K, 1), dtype=np.float64) vhinge_floor_i = np.zeros(max(K, 1), dtype=np.float64) vhinge_cnt_i = np.zeros(max(K, 1), dtype=np.int64) vklq_sum_i = np.zeros(max(K, 1), dtype=np.float64) vheps_sum_i = np.zeros(max(K, 1), dtype=np.float64) # Nesting on EVERY step 1->2, ..., (s-1)->s. A puzzle is prefix-clean # iff no empty cell's nucleus grows on any of those steps. Vacuous at # stage 1 (one latent, no step). Per-step counts are kept so a single # bad transition is visible and not washed out by later clean ones. nest_ok = 0 nest_tot = 0 nest_step_ok = np.zeros(max(K, 1), dtype=np.int64) nest_step_tot = np.zeros(max(K, 1), dtype=np.int64) # Mean nucleus size |A_d| and ground-truth |S_d| over empty cells. nest_asz = np.zeros(max(K, 1), dtype=np.float64) nest_ssz = np.zeros(max(K, 1), dtype=np.float64) nest_scnt = np.zeros(max(K, 1), dtype=np.int64) # Soundness: does A_d still contain the puzzle's unique solution digit? # Diagnostic only -- it reads the answer, so gating on it would be # training a solver rather than the candidate distribution. What makes it # worth logging is that it upper-bounds the achievable solve rate: once a # latent has dropped the true digit, no deeper latent can recover it. keep_ok = 0 # puzzles keeping the truth in EVERY cell, EVERY depth keep_tot = 0 keep_cell_ok = 0 # (cell, depth) pairs keeping it keep_cell_tot = 0 keep_d_ok = np.zeros(max(K, 1), dtype=np.int64) # puzzles, per depth keep_d_tot = np.zeros(max(K, 1), dtype=np.int64) nest_tau = float(getattr(config, "promote_nest_tau", 0.99)) gate_stage = (int(curriculum.stage) if curriculum is not None else 0) # 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) # (bs, R, seq_len) instance-valued copies of input_seq; all zeros on # runs without an instance pool, in which case the train-matched # CE pass is skipped and only the legacy keys are reported. inst_seqs = (np.array(batch_tuple[7]) if len(batch_tuple) > 7 else None) if inst_seqs is not None and not inst_seqs.any(): inst_seqs = 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))) noise_on = bool(int(getattr(config, "latent_noise", 0))) mu_w = mu_b = None if noise_on: mu_w = np.asarray(state.params["latent_mu_w"][0]) mu_b = np.asarray(state.params["latent_mu_b"][0]) 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) z_j = hidden[bidx, si3 - 1 + j] if noise_on: z_j = _latent_mean_slot(mu_w, mu_b, z_j, j) latent_vals[:, j] = z_j # ---- 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], eps=promote_eps) 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["puz_ok"] vceq_grp[s] += qs["puz_count"] vceq_pop[s] += qs["ceq_pop"] vhq_pop[s] += qs["hq_pop"] vlogm_pop[s] += qs["logm_pop"] vpop_cnt[s] += qs["pop_count"] vhinge_sum[s] += qs["hinge"] vhinge_floor[s] += qs["hinge_floor"] vklq_sum[s] += qs["klq"] vheps_sum[s] += qs["hinge_eps"] vhinge_cnt[s] += qs["hinge_cnt"] # ---- The gate's own CE(q||p), on the footing training uses ---- # Two corrections to the pass above, both of which inflate it: # prefix: it reads p off the solution-valued sequence, so every # earlier empty cell hands the model the true digit and the # current cell is settled by elimination, then charged against # a q that conditions on no such thing. Instance-valued # sequences are what training feeds; p is averaged over R of # them to recover a marginal comparable to q. # depth: it runs all K slots (k_budget = K above) and compares # that one full-depth prediction against every stage's q. The # stage being gated is only training `stage` of them. if (K > 0 and cand_targets is not None and q_counts is not None and inst_seqs is not None and gate_stage): # The q being scored is the stage's, but the depth p is read at # is the depth the stage TRAINS at: num_passes = min(pps*stage, # K), the same formula train_and_evaluate uses. These coincide # at pps=1; at pps=K every stage trains at full depth and the # gate must read there too, or it scores a depth the weights # were never fit at. s0 = min(gate_stage, K) - 1 pps = int(getattr(config, "passes_per_stage", 1)) depth_gate = min(max(pps * gate_stage, 1), K) R = int(inst_seqs.shape[1]) act_gate = np.arange(K)[None, :] < depth_gate act_gate = np.broadcast_to(act_gate, (bs, K)) lv_gate = np.zeros((bs, K, config.emb_dim), dtype=np.float32) if recurrent: for j in range(depth_gate): act_j = act_gate & (np.arange(K)[None, :] < j) _, hidden = run_model(build_seq_masked, lv_gate, act_j) z_j = hidden[bidx, si3 - 1 + j] # Training injects mu(h) + sigma*eps, so the gate must # read the mean path; the raw hidden is a latent the # weights were never fit at. if noise_on: z_j = _latent_mean_slot(mu_w, mu_b, z_j, j) lv_gate[:, j] = z_j def _mix_p(depth): act = np.broadcast_to( np.arange(K)[None, :] < depth, (bs, K)) acc = None for r in range(R): lg, _ = run_model( np.asarray(inst_seqs[:, r]), lv_gate, act) sh = lg - lg.max(axis=-1, keepdims=True) pr = np.exp(sh) pr /= pr.sum(axis=-1, keepdims=True) acc = pr if acc is None else acc + pr return acc / max(R, 1) # One readout per active depth so every latent is scored, not # just the newest pair. p_at[d] is the mix with d slots on. p_at = {d: _mix_p(d) for d in range(1, depth_gate + 1)} p_mix = p_at[depth_gate] logp_mix = np.log(np.maximum(p_mix, 1e-30)) 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 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 cells = (input_seq[j, base + 3 * t] * 9 + input_seq[j, base + 3 * t + 1]).astype(np.int64) ok_cell = (cells >= 0) & (cells < 81) bits = cand_targets[j, s0, np.where(ok_cell, cells, 0)] sel = ok_cell & (bits > 0) if not sel.any(): continue qs = _inst_ce_stats( logp_mix[j, v_pos - 1, 1:10][sel], bits[sel], q_counts[j, s0, np.where(ok_cell, cells, 0)][sel], eps=promote_eps) vceq_pop_i[s0] += qs["ceq_pop"] vhq_pop_i[s0] += qs["hq_pop"] vlogm_pop_i[s0] += qs["logm_pop"] vpop_cnt_i[s0] += qs["pop_count"] vhinge_sum_i[s0] += qs["hinge"] vhinge_floor_i[s0] += qs["hinge_floor"] vklq_sum_i[s0] += qs["klq"] vheps_sum_i[s0] += qs["hinge_eps"] vhinge_cnt_i[s0] += qs["hinge_cnt"] empty = [ti for ti in range(len(v_pos)) if ok_cell[ti] and int(bits[ti]) > 0] cells_i = cells keep_puz = True keep_at_d = np.ones(depth_gate, bool) for ti in empty: vp = int(v_pos[ti]) cell = int(cells_i[ti]) true_d = int(puzzle_sol[j][cell]) for d in range(1, depth_gate + 1): pv = p_at[d][j, vp - 1, 1:10] k = _nucleus_size(pv, nest_tau) nest_asz[d - 1] += k sb = int(cand_targets[j, d - 1, cell]) nest_ssz[d - 1] += bin(sb).count("1") if sb > 0 else 0 nest_scnt[d - 1] += 1 if 1 <= true_d <= 9: # The truth survives iff it is among the top-k # digits the nucleus kept. order = np.argsort(-np.asarray(pv)) held = bool((order[:k] == (true_d - 1)).any()) keep_cell_tot += 1 keep_cell_ok += int(held) if not held: keep_at_d[d - 1] = False keep_puz = False if empty: keep_tot += 1 keep_ok += int(keep_puz) for d in range(depth_gate): keep_d_tot[d] += 1 keep_d_ok[d] += int(keep_at_d[d]) if gate_stage < 2: continue # Prefix nesting: one growing cell on any step 1->2, ..., # (s-1)->s spoils the puzzle. Each step is also counted # on its own so a weak early latent stays visible. grew_any = False for d in range(1, depth_gate): grew = False for ti in empty: vp = v_pos[ti] a0 = _nucleus_size( p_at[d][j, vp - 1, 1:10], nest_tau) a1 = _nucleus_size( p_at[d + 1][j, vp - 1, 1:10], nest_tau) if a1 > a0: grew = True break nest_step_tot[d - 1] += 1 nest_step_ok[d - 1] += int(not grew) grew_any = grew_any or grew nest_tot += 1 nest_ok += int(not grew_any) 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) # THE PROMOTION GATE, as a stage average over every |S|>=2 cell: # ceq_pop <= logm_pop <=> ceq_margin >= 0 # This is the stage-level reading of "the value cross-entropy is below the # uniform-candidate entropy". ceq_ok above is the stricter per-cell rate. per_stage_ceq_pop = _per_stage(vceq_pop, vpop_cnt) per_stage_logm_pop = _per_stage(vlogm_pop, vpop_cnt) eval_metrics["per_stage_val_ceq_pop"] = per_stage_ceq_pop eval_metrics["per_stage_val_hq_pop"] = _per_stage(vhq_pop, vpop_cnt) eval_metrics["per_stage_val_logm_pop"] = per_stage_logm_pop eval_metrics["per_stage_val_ceq_margin"] = { s: (per_stage_logm_pop[s] - per_stage_ceq_pop[s] if per_stage_ceq_pop[s] >= 0 else -99.0) for s in per_stage_ceq_pop} # Per-puzzle hinge: mean over puzzles of max(0, CE_i - log|S|_i). Zero iff no # puzzle is still above its own bar, so unlike ceq_margin it cannot be passed # by cross-puzzle cancellation. `hinge_floor` is the same statistic at p == q # and bounds how low `hinge` can ever go on this data. eval_metrics["per_stage_val_hinge"] = _per_stage(vhinge_sum, vhinge_cnt) eval_metrics["per_stage_val_hinge_floor"] = _per_stage( vhinge_floor, vhinge_cnt) # Train-matched readout: only the gate stage is populated, the rest read # -1 and the promotion logic falls back to the legacy key for those. eval_metrics["per_stage_val_hinge_inst"] = _per_stage( vhinge_sum_i, vhinge_cnt_i) eval_metrics["per_stage_val_hinge_floor_inst"] = _per_stage( vhinge_floor_i, vhinge_cnt_i) eval_metrics["per_stage_val_klq"] = _per_stage(vklq_sum, vhinge_cnt) eval_metrics["per_stage_val_klq_inst"] = _per_stage(vklq_sum_i, vhinge_cnt_i) eval_metrics["per_stage_val_hinge_eps"] = _per_stage(vheps_sum, vhinge_cnt) eval_metrics["per_stage_val_hinge_eps_inst"] = _per_stage( vheps_sum_i, vhinge_cnt_i) per_stage_ceq_pop_i = _per_stage(vceq_pop_i, vpop_cnt_i) per_stage_logm_pop_i = _per_stage(vlogm_pop_i, vpop_cnt_i) eval_metrics["per_stage_val_ceq_pop_inst"] = per_stage_ceq_pop_i eval_metrics["per_stage_val_hq_pop_inst"] = _per_stage( vhq_pop_i, vpop_cnt_i) eval_metrics["per_stage_val_ceq_margin_inst"] = { s: (per_stage_logm_pop_i[s] - per_stage_ceq_pop_i[s] if per_stage_ceq_pop_i[s] >= 0 else -99.0) for s in per_stage_ceq_pop_i} 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) # Puzzle-level nesting over the whole prefix 1->2->...->s. Empty at # stage 1 so the printer reports -1 and the gate treats it as vacuous. # nest_clean = fraction of puzzles with no growing cell on ANY step. # nest_step = per-step fraction (keyed by the later depth: 2 means 1->2). if nest_tot: eval_metrics["nest_clean"] = [float(nest_ok) / float(nest_tot)] eval_metrics["nest_n"] = [float(nest_tot)] eval_metrics["nest_grew"] = [1.0 - float(nest_ok) / float(nest_tot)] eval_metrics["nest_step"] = { d + 2: (float(nest_step_ok[d] / nest_step_tot[d]) if nest_step_tot[d] else -1.0) for d in range(K - 1) } step_fracs = [eval_metrics["nest_step"][s] for s in range(2, gate_stage + 1) if eval_metrics["nest_step"].get(s, -1.0) >= 0] eval_metrics["nest_min"] = [min(step_fracs)] if step_fracs else [] if nest_scnt[0] > 0: eval_metrics["nest_A"] = { d + 1: (float(nest_asz[d] / nest_scnt[d]) if nest_scnt[d] else -1.0) for d in range(K) } eval_metrics["nest_S"] = { d + 1: (float(nest_ssz[d] / nest_scnt[d]) if nest_scnt[d] else -1.0) for d in range(K) } eval_metrics["nest_A1"] = [eval_metrics["nest_A"][1]] eval_metrics["nest_S1"] = [eval_metrics["nest_S"][1]] if keep_tot: # keep_all = puzzles where no cell lost the true digit at any depth, # i.e. the fraction still solvable in principle. # keep_cell = the same over (cell, depth), which stays high even when # keep_all collapses, since one bad cell kills a puzzle. eval_metrics["keep_all"] = [float(keep_ok) / float(keep_tot)] eval_metrics["keep_cell"] = [float(keep_cell_ok) / float(keep_cell_tot)] eval_metrics["keep_depth"] = { d + 1: (float(keep_d_ok[d] / keep_d_tot[d]) if keep_d_tot[d] else -1.0) for d in range(K) } return eval_metrics