"""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