"""Evaluation related functions.""" 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 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": [], # Value/placement acc: correct digit at the model-chosen cell "loc_acc": [], # Location acc: model picks the ground-truth next cell (r,c) "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-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)) 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)) 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:]]) # ---- 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'}") 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}}}") 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 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 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 if 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 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 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