Add README and training/generation code

README.md

@@ -0,0 +1,72 @@
+---
+license: apache-2.0
+tags:
+  - sudoku
+  - reasoning
+  - curriculum
+  - jax
+---
+
+# Sudoku superposition instances
+
+Concrete board assignments for a 12-stage Sudoku latent curriculum.
+Each stage keeps a candidate set per cell; this dataset materializes
+those sets as ordinary `(row, col, value)` sequences so training is
+standard next-token cross-entropy (no multi-hot candidate head).
+
+Repo: [Avra98/Sudoku_superposition](https://huggingface.co/Avra98/Sudoku_superposition)
+
+## Dataset
+
+| split | puzzles | instances | mean / puzzle |
+| --- | ---: | ---: | ---: |
+| train | 1,804,462 | 89,236,838 | 49.45 |
+| test | 99,999 | 4,945,532 | 49.46 |
+
+Mean instances per stage (train, stage 0 → 11):
+`6.28, 5.96, 5.65, 5.34, 5.02, 4.68, 4.32, 3.88, 3.33, 2.58, 1.41, 1.00`
+
+Stage 11 is the unique solution. No empty `(puzzle, stage)` pair.
+
+### Files (`data/`)
+
+| file | shape | dtype | role |
+| --- | --- | --- | --- |
+| `{split}_assignments.npy` | `(M, 81)` | uint8 | one full board per instance, cell `r*9+c` |
+| `{split}_starts.npy` | `(N, 12)` | int32 | first row in `assignments` for `(puzzle, stage)` |
+| `{split}_counts.npy` | `(N, 12)` | uint8 | number of instances for `(puzzle, stage)` |
+| `{split}_index.npy` | `(M, 3)` | int32 | `[puzzle_idx, stage, k]` (optional; starts/counts are enough) |
+
+The trainer only needs `assignments`, `starts`, and `counts`.
+
+Puzzle clue/solution arrays are **not** in this repo (they are the
+original Sudoku npy files). Candidate masks used to *build* the
+instances live in `datasets_multicandidate_s12/`.
+
+### How a training example is built
+
+1. Take the usual solver-order sequence: clue triples, then K latent
+   placeholders, then empty-cell triples.
+2. Sample one instance for the current curriculum stage.
+3. Rewrite **values only**. Location order stays solver-order.
+4. Predict the output triples with softmax CE. Latent slots are not
+   predicted.
+
+Curriculum stage `t` (1..12) trains on stage-`(t-1)` instances.
+Stage 12 targets the unique solution.
+
+## Code (`code/`)
+
+- `code/train/` — JAX trainer (`data.py`, `trainer.py`, `evaluater.py`,
+  `train_and_evaluate.py`, `train_backtrack.py`, `main.py`, `model.py`)
+- `code/build_superposition_dataset.py` — instance generator
+- `code/build_instance_offsets.py` — `starts` / `counts` tables
+- `code/superposition_instances.py` — per-puzzle instance sampler
+- `code/sbatch_instance_latent.sh` — Slurm launch (feanor / H200)
+
+Set `SUDOKU_INSTANCE_DIR` to the `data/` folder (or a local copy).
+
+```bash
+from huggingface_hub import snapshot_download
+snapshot_download("Avra98/Sudoku_superposition", local_dir="Sudoku_superposition")
+```

code/build_instance_offsets.py

@@ -0,0 +1,46 @@
+"""Precompute (puzzle, stage) -> row range into the instance assignments.
+
+The index rows are grouped by puzzle then stage, so a start offset plus a count
+per (puzzle, stage) is enough for O(1) sampling in the data loader.
+
+Writes {split}_starts.npy (N, S) int32 and {split}_counts.npy (N, S) uint8.
+"""
+import argparse
+import os
+
+import numpy as np
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--dir", default="datasets_superposition")
+    ap.add_argument("--splits", default="train,test")
+    ap.add_argument("--stages", type=int, default=12)
+    args = ap.parse_args()
+
+    S = args.stages
+    for split in args.splits.split(","):
+        split = split.strip()
+        ix = np.load(os.path.join(args.dir, f"{split}_index.npy"), mmap_mode="r")
+        n_puzzles = int(np.array(ix[-1, 0])) + 1
+        key = np.array(ix[:, 0]).astype(np.int64) * S + np.array(ix[:, 1])
+        counts = np.bincount(key, minlength=n_puzzles * S).astype(np.int64)
+        if counts.max() > 255:
+            raise ValueError(f"count {counts.max()} exceeds uint8")
+        starts = np.concatenate([[0], np.cumsum(counts)[:-1]]).astype(np.int64)
+
+        starts = starts.reshape(n_puzzles, S).astype(np.int32)
+        counts = counts.reshape(n_puzzles, S).astype(np.uint8)
+
+        np.save(os.path.join(args.dir, f"{split}_starts.npy"), starts)
+        np.save(os.path.join(args.dir, f"{split}_counts.npy"), counts)
+
+        empty = int((counts == 0).sum())
+        print(f"[{split}] puzzles={n_puzzles:,} rows={len(ix):,}  "
+              f"(puzzle,stage) cells with zero instances: {empty:,}")
+        print(f"    mean instances per stage: "
+              + " ".join(f"{counts[:, s].mean():.2f}" for s in range(S)))
+
+
+if __name__ == "__main__":
+    main()

code/build_superposition_dataset.py

@@ -0,0 +1,128 @@
+"""Materialize superposition instances for every puzzle and every stage.
+
+Input:  datasets_multicandidate_s12/{split}_cand_masks.npy   (N, 12, 81)
+Output: datasets_superposition/{split}_assignments.npy       (M, 81) uint8
+        datasets_superposition/{split}_index.npy             (M, 3)  int32
+            columns: [puzzle_idx, stage, instance_within_stage]
+
+M is about 48 x N (one puzzle expands into ~6 instances at stage 0 down to 1
+at stage 11). Generation is embarrassingly parallel over puzzles.
+"""
+import argparse
+import os
+import time
+
+import numpy as np
+import multiprocessing as mp
+
+import superposition_instances as SI
+
+_MASKS = None
+_ARGS = None
+
+
+def _init(path, args):
+    global _MASKS, _ARGS
+    _MASKS = np.load(path, mmap_mode="r")
+    _ARGS = args
+
+
+def _one(idx):
+    assigns, index = [], []
+    for s in range(_MASKS.shape[1]):
+        mask = np.array(_MASKS[idx, s]).astype(np.uint16)
+        r = SI.instances_for_stage(
+            mask,
+            max_confine=_ARGS.max_confine,
+            max_instances=_ARGS.max_instances,
+            max_attempts=_ARGS.max_attempts,
+            seed=idx * 100 + s,
+            max_repair=_ARGS.max_repair,
+            require_new=True,
+            patience=_ARGS.patience,
+        )
+        inst = r["instances"]
+        if len(inst) == 0:
+            continue
+        assigns.append(inst.astype(np.uint8))
+        for k in range(len(inst)):
+            index.append((idx, s, k))
+    if not assigns:
+        return (np.zeros((0, 81), dtype=np.uint8),
+                np.zeros((0, 3), dtype=np.int32))
+    return (np.concatenate(assigns, axis=0),
+            np.array(index, dtype=np.int32))
+
+
+def process_split(split, args):
+    mask_path = os.path.join(args.mask_dir, f"{split}_cand_masks.npy")
+    masks = np.load(mask_path, mmap_mode="r")
+    n = len(masks) if args.limit is None else min(args.limit, len(masks))
+    print(f"[{split}] {n:,} puzzles from {mask_path}", flush=True)
+
+    t0 = time.time()
+    chunks_a, chunks_i = [], []
+    done = 0
+    with mp.Pool(args.workers, initializer=_init,
+                 initargs=(mask_path, args)) as pool:
+        for a, i in pool.imap(_one, range(n), chunksize=8):
+            if len(a):
+                chunks_a.append(a)
+                chunks_i.append(i)
+            done += 1
+            if done % 2000 == 0 or done == n:
+                rate = done / max(time.time() - t0, 1e-6)
+                kept = sum(len(x) for x in chunks_i)
+                print(f"  {done:,}/{n:,}  {rate:.0f} puzzles/s  "
+                      f"{kept:,} instances  "
+                      f"({kept / done:.1f} per puzzle)", flush=True)
+
+    assignments = (np.concatenate(chunks_a, axis=0) if chunks_a
+                   else np.zeros((0, 81), dtype=np.uint8))
+    index = (np.concatenate(chunks_i, axis=0) if chunks_i
+             else np.zeros((0, 3), dtype=np.int32))
+
+    os.makedirs(args.out_dir, exist_ok=True)
+    ap = os.path.join(args.out_dir, f"{split}_assignments.npy")
+    ip = os.path.join(args.out_dir, f"{split}_index.npy")
+    np.save(ap, assignments)
+    np.save(ip, index)
+
+    elapsed = time.time() - t0
+    print(f"[{split}] saved {len(assignments):,} instances  "
+          f"({len(assignments) / n:.2f} per puzzle) in {elapsed / 60:.1f} min",
+          flush=True)
+    print(f"         {ap}  {os.path.getsize(ap) / 1e9:.2f} GB", flush=True)
+    print(f"         {ip}  {os.path.getsize(ip) / 1e9:.2f} GB", flush=True)
+
+    # per-stage instance counts
+    print(f"[{split}] instances per stage:", flush=True)
+    for s in range(masks.shape[1]):
+        c = int((index[:, 1] == s).sum())
+        print(f"    stage {s:>2}: {c:>12,}  ({c / n:.2f} per puzzle)",
+              flush=True)
+    return len(assignments)
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--mask_dir", default="datasets_multicandidate_s12")
+    ap.add_argument("--out_dir", default="datasets_superposition")
+    ap.add_argument("--splits", default="train,test")
+    ap.add_argument("--limit", type=int, default=None)
+    ap.add_argument("--max-confine", type=int, default=1)
+    ap.add_argument("--max-instances", type=int, default=32)
+    ap.add_argument("--max-attempts", type=int, default=400)
+    ap.add_argument("--max-repair", type=int, default=80)
+    ap.add_argument("--patience", type=int, default=40)
+    ap.add_argument("--workers", type=int, default=64)
+    args = ap.parse_args()
+
+    print(f"confine |S|<={args.max_confine}  workers={args.workers}  "
+          f"out={args.out_dir}", flush=True)
+    for split in args.splits.split(","):
+        process_split(split.strip(), args)
+
+
+if __name__ == "__main__":
+    main()

code/env_paths.sh

@@ -0,0 +1,77 @@
+#!/usr/bin/env bash
+# Portable RUN_DIR / PY / CUDNN_LIB for slimgpu vs Berkeley (gandalf/feanor).
+#
+# Override any of these before sourcing:
+#   SUDOKU_RUN_DIR, SUDOKU_PY, SUDOKU_CUDNN_LIB
+#
+# Usage (from launch_*.sh):
+#   source "$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)/env_paths.sh"
+
+_SLIM_ROOT="/egr/research-slim/ghoshavr/llm-reasoning-logic-puzzles"
+_SLIM_RUN="${_SLIM_ROOT}/sudoku-code/multicandidate_run"
+_SLIM_PY="/egr/research-slim/ghoshavr/conda-envs/logicpuzzles/bin/python"
+_SLIM_CUDNN="/egr/research-slim/ghoshavr/conda-envs/logicpuzzles/lib/python3.9/site-packages/nvidia/cudnn/lib"
+
+# Berkeley: prefer scratch (home often quota-limited), then $HOME (may be a symlink).
+_BERK_RUN_SCRATCH="/scratch/users/gatmiry/llm-reasoning-logic-puzzles/sudoku-code/multicandidate_run"
+_BERK_RUN="${HOME}/llm-reasoning-logic-puzzles/sudoku-code/multicandidate_run"
+
+# Prefer explicit overrides, then slimgpu tree if present, else Berkeley home.
+if [ -n "${SUDOKU_RUN_DIR:-}" ]; then
+  RUN_DIR="${SUDOKU_RUN_DIR}"
+elif [ -d "${_SLIM_RUN}" ]; then
+  RUN_DIR="${_SLIM_RUN}"
+elif [ -d "${_BERK_RUN_SCRATCH}" ]; then
+  RUN_DIR="${_BERK_RUN_SCRATCH}"
+else
+  RUN_DIR="${_BERK_RUN}"
+fi
+
+if [ -n "${SUDOKU_PY:-}" ]; then
+  PY="${SUDOKU_PY}"
+elif [ -x "${_SLIM_PY}" ]; then
+  PY="${_SLIM_PY}"
+else
+  # Common Berkeley conda layouts; first existing wins.
+  for _cand in \
+      "${HOME}/miniconda3/envs/logic_puzzles/bin/python" \
+      "${HOME}/miniconda3/envs/logicpuzzles/bin/python" \
+      "${HOME}/anaconda3/envs/logic_puzzles/bin/python" \
+      "${HOME}/anaconda3/envs/logicpuzzles/bin/python" \
+      "${HOME}/.conda/envs/logic_puzzles/bin/python" \
+      "${HOME}/.conda/envs/logicpuzzles/bin/python" \
+      "/scratch/users/gatmiry/conda/envs/logic_puzzles/bin/python" \
+      "/scratch/users/gatmiry/conda/envs/logicpuzzles/bin/python"; do
+    if [ -x "${_cand}" ]; then
+      PY="${_cand}"
+      break
+    fi
+  done
+  if [ -z "${PY:-}" ]; then
+    if command -v python >/dev/null 2>&1; then
+      PY="$(command -v python)"
+    else
+      echo "ERROR: no Python found. Create env: conda env create -f ~/llm-reasoning-logic-puzzles/environment.yml -n logic_puzzles" >&2
+      echo "  then: export SUDOKU_PY=\$HOME/miniconda3/envs/logic_puzzles/bin/python" >&2
+      return 1 2>/dev/null || exit 1
+    fi
+  fi
+fi
+
+if [ -n "${SUDOKU_CUDNN_LIB:-}" ]; then
+  CUDNN_LIB="${SUDOKU_CUDNN_LIB}"
+elif [ -d "${_SLIM_CUDNN}" ]; then
+  CUDNN_LIB="${_SLIM_CUDNN}"
+else
+  # Derive from PY site-packages if present; else leave empty (cluster CUDA often enough).
+  _py_root="$(cd "$(dirname "${PY}")/.." && pwd)"
+  _cudnn_guess="${_py_root}/lib/python3.9/site-packages/nvidia/cudnn/lib"
+  if [ -d "${_cudnn_guess}" ]; then
+    CUDNN_LIB="${_cudnn_guess}"
+  else
+    CUDNN_LIB=""
+  fi
+fi
+
+export RUN_DIR PY CUDNN_LIB
+unset _SLIM_ROOT _SLIM_RUN _SLIM_PY _SLIM_CUDNN _BERK_RUN _cand _py_root _cudnn_guess

code/launch_instance_latent.sh

@@ -0,0 +1,73 @@
+#!/usr/bin/env bash
+# 12-stage latent curriculum, instance-wise superposition targets.
+#
+# Same architecture and curriculum as w12_latent, except:
+#   - the output prompt is one sampled stage-k assignment, not the unique
+#     solution (input clues stay fixed)
+#   - the candidate-head BCE is off (SUDOKU_AUX_WEIGHT=0)
+#   - promotion is gated on the in-set rate (emitted digit is a stage-k
+#     candidate), not on candidate-set exact match
+#
+# Usage:
+#   bash launch_instance_latent.sh            # GPU 0, full 800k run
+#   GPU=6 bash launch_instance_latent.sh      # pick another GPU
+#   SUDOKU_MAX_STEPS=200 bash launch_instance_latent.sh   # smoke
+
+set -u
+
+_SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
+_CODE_DIR="$(cd "${_SCRIPT_DIR}/.." && pwd)"
+export SUDOKU_RUN_DIR="${SUDOKU_RUN_DIR:-${_SCRIPT_DIR}}"
+# shellcheck source=env_paths.sh
+source "${_SCRIPT_DIR}/env_paths.sh"
+
+cd "$RUN_DIR" || exit 1
+if [ -n "${CUDNN_LIB}" ]; then export LD_LIBRARY_PATH="${CUDNN_LIB}:${LD_LIBRARY_PATH:-}"; fi
+export XLA_PYTHON_CLIENT_MEM_FRACTION=0.9
+
+GPU="${GPU:-0}"
+name="${NAME:-w12_inst_latent}"
+mkdir -p logs
+
+export SUDOKU_TRAIN_PATH="${_CODE_DIR}/datasets/train_sudoku_puzzles.npy"
+export SUDOKU_TEST_PATH="${_CODE_DIR}/datasets/test_sudoku_puzzles.npy"
+# Masks are a metric only (in-set rate / promotion). Not a training target.
+export SUDOKU_TRAIN_CAND="${_CODE_DIR}/datasets_multicandidate_s12/train_cand_masks.npy"
+export SUDOKU_TEST_CAND="${_CODE_DIR}/datasets_multicandidate_s12/test_cand_masks.npy"
+export SUDOKU_INSTANCE_DIR="${_CODE_DIR}/datasets_superposition"
+
+export SUDOKU_RESUME=0
+export SUDOKU_LATENT_SLOTS=12
+export SUDOKU_RECURRENT=1
+export SUDOKU_BACKTRACK=0
+export SUDOKU_START_STAGE=1
+export SUDOKU_MAX_STAGE=12
+export SUDOKU_CAND_SLOT_MODE=depth
+export SUDOKU_PASSES_PER_STAGE=1
+export SUDOKU_AUX_WEIGHT=0.0
+export SUDOKU_LEVEL_BALANCED=0
+export SUDOKU_DATA_CURRICULUM=none
+
+export SUDOKU_PLATEAU_STEPS="${SUDOKU_PLATEAU_STEPS:-20000}"
+export SUDOKU_PLATEAU_DELTA=0.005
+export SUDOKU_PATIENCE="${SUDOKU_PATIENCE:-80000}"
+export SUDOKU_MIN_STAGE_STEPS="${SUDOKU_MIN_STAGE_STEPS:-8000}"
+export SUDOKU_PROMOTE_ACC="${SUDOKU_PROMOTE_ACC:-0.90}"
+export SUDOKU_MAX_STEPS="${SUDOKU_MAX_STEPS:-800000}"
+export SUDOKU_EVAL_EVERY="${SUDOKU_EVAL_EVERY:-2000}"
+export SUDOKU_SAVE_EVERY="${SUDOKU_SAVE_EVERY:-10000}"
+export SUDOKU_CKPT_KEEP=100
+
+export SUDOKU_LR=0.0002
+export SUDOKU_DROPOUT=0.2
+export SUDOKU_WD=0.005
+
+echo "launching ${name} on GPU ${GPU}"
+echo "  K=12 recurrent=1 aux=0 instance_dir=${SUDOKU_INSTANCE_DIR}"
+echo "  max_steps=${SUDOKU_MAX_STEPS} plateau=${SUDOKU_PLATEAU_STEPS} patience=${SUDOKU_PATIENCE}"
+CUDA_VISIBLE_DEVICES="${GPU}" \
+  nohup "$PY" -u -m train.main \
+    --workdir="./logs/${name}" \
+    --exp_name="${name}" \
+    > "logs/${name}.log" 2>&1 &
+echo "  pid $!  log=${RUN_DIR}/logs/${name}.log"

code/level_instance_stats.py

@@ -0,0 +1,94 @@
+"""Candidate-set size and instance count per stage, level 3 vs level 8.
+
+Two metrics only, measured on the training masks.
+"""
+import argparse
+import multiprocessing as mp
+
+import numpy as np
+
+import superposition_instances as SI
+
+_MASKS = None
+_ARGS = None
+
+
+def _init(path, args):
+    global _MASKS, _ARGS
+    _MASKS = np.load(path, mmap_mode="r")
+    _ARGS = args
+
+
+def _one(idx):
+    out = []
+    for s in range(_MASKS.shape[1]):
+        mask = np.array(_MASKS[idx, s]).astype(np.uint16)
+        _, empties = SI.split_cells(mask)
+        sizes = [len(SI.digits_of(mask[c])) for c in empties]
+        width = float(np.mean(sizes)) if sizes else 1.0
+        wmax = int(np.max(sizes)) if sizes else 1
+        r = SI.instances_for_stage(
+            mask, max_confine=_ARGS.max_confine,
+            max_instances=_ARGS.max_instances,
+            max_attempts=_ARGS.max_attempts, seed=idx * 100 + s,
+            max_repair=_ARGS.max_repair)
+        out.append((s, width, wmax, r["n_instances"], len(empties)))
+    return out
+
+
+def run(level, args):
+    meta = np.load(args.meta)
+    idxs = np.where(meta[:, 1] == level)[0][:args.limit].tolist()
+    with mp.Pool(args.workers, initializer=_init,
+                 initargs=(args.masks, args)) as pool:
+        res = pool.map(_one, idxs, chunksize=4)
+    widths, maxes, counts, opens = {}, {}, {}, {}
+    for rows in res:
+        for s, w, wm, n, ne in rows:
+            widths.setdefault(s, []).append(w)
+            maxes.setdefault(s, []).append(wm)
+            counts.setdefault(s, []).append(n)
+            opens.setdefault(s, []).append(ne)
+    return len(idxs), widths, maxes, counts, opens
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--masks", default="datasets_multicandidate_s12/train_cand_masks.npy")
+    ap.add_argument("--meta", default="datasets_multicandidate_s12/train_meta.npy")
+    ap.add_argument("--limit", type=int, default=400)
+    ap.add_argument("--max-confine", type=int, default=1)
+    ap.add_argument("--max-instances", type=int, default=64)
+    ap.add_argument("--max-attempts", type=int, default=300)
+    ap.add_argument("--max-repair", type=int, default=120)
+    ap.add_argument("--workers", type=int, default=60)
+    args = ap.parse_args()
+
+    n3, w3, m3, c3, o3 = run(3, args)
+    n8, w8, m8, c8, o8 = run(8, args)
+
+    print(f"training masks, {n3} level-3 puzzles and {n8} level-8 puzzles")
+    print(f"confinement threshold |S| <= {args.max_confine}")
+    print()
+    print(f"{'':>5} {'------------- level 3 -------------':>42} "
+          f"{'------------- level 8 -------------':>42}")
+    print(f"{'stage':>5} {'mean':>7} {'max':>7} {'worst':>7} {'open':>7} "
+          f"{'inst':>7} {'mean':>7} {'max':>7} {'worst':>7} {'open':>7} {'inst':>7}")
+    print("-" * 96)
+    for s in sorted(w3):
+        print(f"{s:>5} "
+              f"{np.mean(w3[s]):>7.2f} {np.mean(m3[s]):>7.2f} "
+              f"{np.max(m3[s]):>7d} {np.mean(o3[s]):>7.1f} "
+              f"{np.mean(c3[s]):>7.1f} "
+              f"{np.mean(w8[s]):>7.2f} {np.mean(m8[s]):>7.2f} "
+              f"{np.max(m8[s]):>7d} {np.mean(o8[s]):>7.1f} "
+              f"{np.mean(c8[s]):>7.1f}")
+    print()
+    print("mean  = mean candidate set size over undetermined cells")
+    print("max   = mean over puzzles of the largest candidate set in that puzzle")
+    print("worst = largest candidate set seen in any puzzle at that stage")
+    print("open  = undetermined cells remaining;  inst = instances at saturation")
+
+
+if __name__ == "__main__":
+    main()

code/sbatch_instance_latent.sh

@@ -0,0 +1,142 @@
+#!/bin/bash
+#SBATCH --partition=songmei
+#SBATCH --nodelist=feanor
+#SBATCH --gres=gpu:1
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=64G
+#SBATCH --time=72:00:00
+#SBATCH --job-name=w12_inst
+#SBATCH --output=/scratch/users/gatmiry/llm-reasoning-logic-puzzles/sudoku-code/wavecurriculum_run/logs/slurm_%x_%j.out
+#SBATCH --error=/scratch/users/gatmiry/llm-reasoning-logic-puzzles/sudoku-code/wavecurriculum_run/logs/slurm_%x_%j.err
+
+# Same recipe as the slimgpu w12_inst_latent run:
+#   K=12 recurrent latent curriculum, no backtrack, no candidate-head BCE.
+#   Output prompt = one sampled stage-k instance (CE against that sequence).
+# Instance npy files and s12 masks live on feanor /tmp (scratch quota is 20 G).
+
+set -u
+hostname
+nvidia-smi -L
+echo "[$(date)] w12_inst_latent"
+
+SCRATCH_ROOT=/scratch/users/gatmiry/llm-reasoning-logic-puzzles
+RUN_DIR=${SCRATCH_ROOT}/sudoku-code/wavecurriculum_run
+ENV_LOCAL=/tmp/logicpuzzles
+CAND_DIR=/tmp/sudoku_s12
+INST_DIR=/tmp/sudoku_superposition
+LOCAL_LOG=/tmp/sudoku_wave_runs/w12_inst_latent
+TARBALL_GANDALF=/tmp/logicpuzzles_env.tar.gz
+TARBALL_LOCAL=/tmp/logicpuzzles_env_${SLURM_JOB_ID}.tar.gz
+GANDALF_INST=gandalf.berkeley.edu:/tmp/sudoku_superposition
+GANDALF_CAND=gandalf.berkeley.edu:/tmp/sudoku_s12
+
+mkdir -p "${RUN_DIR}/logs" "${LOCAL_LOG}" /tmp/sudoku_wave_runs
+
+SCP_OPTS="-o IdentitiesOnly=yes -o StrictHostKeyChecking=accept-new"
+[ -f "${HOME}/.ssh/id_ed25519_berkeley" ] && SCP_OPTS="${SCP_OPTS} -i ${HOME}/.ssh/id_ed25519_berkeley"
+
+if ${ENV_LOCAL}/bin/python -u -c "import jax; assert jax.default_backend()=='gpu' or 'cuda' in str(jax.devices()[0]).lower()" 2>/dev/null; then
+  echo "[$(date)] reusing ${ENV_LOCAL}"
+else
+  echo "[$(date)] fetching env tarball"
+  rm -rf "${ENV_LOCAL}"
+  scp ${SCP_OPTS} "gandalf.berkeley.edu:${TARBALL_GANDALF}" "${TARBALL_LOCAL}"
+  tar xzf "${TARBALL_LOCAL}" -C /tmp
+  rm -f "${TARBALL_LOCAL}"
+fi
+
+export PY=${ENV_LOCAL}/bin/python
+export LD_LIBRARY_PATH=\
+${ENV_LOCAL}/lib/python3.9/site-packages/nvidia/cudnn/lib:\
+${ENV_LOCAL}/lib/python3.9/site-packages/nvidia/cublas/lib:\
+${ENV_LOCAL}/lib/python3.9/site-packages/nvidia/cuda_runtime/lib:\
+${ENV_LOCAL}/lib/python3.9/site-packages/nvidia/cuda_nvrtc/lib:\
+${ENV_LOCAL}/lib/python3.9/site-packages/nvidia/nccl/lib:\
+${LD_LIBRARY_PATH:-}
+
+${PY} -u -c "import jax; print(jax.__version__, jax.devices(), jax.default_backend())"
+
+# ---- Stage instance files onto node-local /tmp (scratch cannot hold 8 G) ----
+need_inst=0
+for f in train_assignments.npy train_starts.npy train_counts.npy \
+         test_assignments.npy test_starts.npy test_counts.npy; do
+  [ -s "${INST_DIR}/${f}" ] || need_inst=1
+done
+if [ "${need_inst}" = 1 ]; then
+  echo "[$(date)] pulling instances from ${GANDALF_INST}"
+  mkdir -p "${INST_DIR}"
+  rsync -a --progress -e "ssh ${SCP_OPTS}" \
+    "${GANDALF_INST}/" "${INST_DIR}/"
+fi
+for f in train_assignments.npy train_starts.npy train_counts.npy \
+         test_assignments.npy test_starts.npy test_counts.npy; do
+  [ -s "${INST_DIR}/${f}" ] || { echo "missing ${INST_DIR}/${f}" >&2; exit 1; }
+done
+
+need_cand=0
+for f in train_cand_masks.npy test_cand_masks.npy; do
+  [ -s "${CAND_DIR}/${f}" ] || need_cand=1
+done
+if [ "${need_cand}" = 1 ]; then
+  echo "[$(date)] pulling s12 masks from ${GANDALF_CAND}"
+  mkdir -p "${CAND_DIR}"
+  rsync -a -e "ssh ${SCP_OPTS}" "${GANDALF_CAND}/" "${CAND_DIR}/" || true
+fi
+for f in train_cand_masks.npy test_cand_masks.npy; do
+  [ -s "${CAND_DIR}/${f}" ] || { echo "missing ${CAND_DIR}/${f}" >&2; exit 1; }
+done
+
+# ---- Exact same recipe as slimgpu w12_inst_latent ----
+export SUDOKU_RESUME=0
+export SUDOKU_START_STAGE=1
+export SUDOKU_MAX_STAGE=12
+export SUDOKU_LATENT_SLOTS=12
+export SUDOKU_RECURRENT=1
+export SUDOKU_BACKTRACK=0
+export SUDOKU_CAND_SLOT_MODE=depth
+export SUDOKU_PASSES_PER_STAGE=1
+export SUDOKU_AUX_WEIGHT=0.0
+export SUDOKU_LEVEL_BALANCED=0
+export SUDOKU_DATA_CURRICULUM=none
+export SUDOKU_PLATEAU_STEPS=20000
+export SUDOKU_PLATEAU_DELTA=0.005
+export SUDOKU_PATIENCE=80000
+export SUDOKU_MIN_STAGE_STEPS=8000
+# Stage 1->2: inset >= 0.85 (chance is ~0.41 on stage-0 sets) AND
+# loc_acc >= 0.70. Frontier is min(inset, loc), so both must clear 0.85
+# unless loc is the limiter (then 0.85 loc is required). The loc floor
+# is applied via min(); SUDOKU_PROMOTE_LOC is the documented intent.
+export SUDOKU_PROMOTE_ACC=0.85
+export SUDOKU_PROMOTE_LOC=0.70
+export SUDOKU_MAX_STEPS="${SUDOKU_MAX_STEPS:-800000}"
+export SUDOKU_EVAL_EVERY=2000
+export SUDOKU_SAVE_EVERY=10000
+export SUDOKU_CKPT_KEEP=3
+export SUDOKU_LR=0.0002
+export SUDOKU_DROPOUT=0.2
+export SUDOKU_WD=0.005
+export SUDOKU_TRAIN_PATH="${SCRATCH_ROOT}/sudoku-code/datasets/train_sudoku_puzzles.npy"
+export SUDOKU_TEST_PATH="${SCRATCH_ROOT}/sudoku-code/datasets/test_sudoku_puzzles.npy"
+export SUDOKU_TRAIN_CAND="${CAND_DIR}/train_cand_masks.npy"
+export SUDOKU_TEST_CAND="${CAND_DIR}/test_cand_masks.npy"
+export SUDOKU_INSTANCE_DIR="${INST_DIR}"
+export XLA_PYTHON_CLIENT_MEM_FRACTION=0.9
+
+cd "${RUN_DIR}"
+(
+  while true; do sleep 900
+    rsync -a "${LOCAL_LOG}.log" "${RUN_DIR}/logs/w12_inst_latent.log" 2>/dev/null || true
+  done
+) &
+SYNC_PID=$!
+
+echo "[$(date)] starting w12_inst_latent from scratch"
+echo "  K=12 recurrent=1 bt=0 aux=0 instance_dir=${INST_DIR}"
+CUDA_VISIBLE_DEVICES=0 ${PY} -u -m train.main \
+  --workdir="${LOCAL_LOG}" --exp_name="w12_inst_latent" \
+  > "${LOCAL_LOG}.log" 2>&1
+EC=$?
+kill ${SYNC_PID} 2>/dev/null || true
+rsync -a "${LOCAL_LOG}.log" "${RUN_DIR}/logs/w12_inst_latent.log" 2>/dev/null || true
+echo "[$(date)] w12_inst_latent exit ${EC}"
+exit ${EC}

code/sbatch_instance_latent_bt.sh

@@ -0,0 +1,145 @@
+#!/bin/bash
+#SBATCH --partition=songmei
+#SBATCH --nodelist=feanor
+#SBATCH --gres=gpu:1
+#SBATCH --cpus-per-task=8
+#SBATCH --mem=64G
+#SBATCH --time=72:00:00
+#SBATCH --job-name=w12_instbt
+#SBATCH --output=/scratch/users/gatmiry/llm-reasoning-logic-puzzles/sudoku-code/wavecurriculum_run/logs/slurm_%x_%j.out
+#SBATCH --error=/scratch/users/gatmiry/llm-reasoning-logic-puzzles/sudoku-code/wavecurriculum_run/logs/slurm_%x_%j.err
+
+# Same instance-CE latent curriculum as w12_inst, plus adaptive backtrack:
+# if a graduated stage's in-set rate falls more than the margin below its
+# graduation value, train at that stage's depth until it recovers.
+
+set -u
+hostname
+nvidia-smi -L
+echo "[$(date)] w12_inst_latent_bt"
+
+SCRATCH_ROOT=/scratch/users/gatmiry/llm-reasoning-logic-puzzles
+RUN_DIR=${SCRATCH_ROOT}/sudoku-code/wavecurriculum_run
+ENV_LOCAL=/tmp/logicpuzzles
+CAND_DIR=/tmp/sudoku_s12
+INST_DIR=/tmp/sudoku_superposition
+LOCAL_LOG=/tmp/sudoku_wave_runs/w12_inst_latent_bt
+TARBALL_GANDALF=/tmp/logicpuzzles_env.tar.gz
+TARBALL_LOCAL=/tmp/logicpuzzles_env_${SLURM_JOB_ID}.tar.gz
+GANDALF_INST=gandalf.berkeley.edu:/tmp/sudoku_superposition
+GANDALF_CAND=gandalf.berkeley.edu:/tmp/sudoku_s12
+
+mkdir -p "${RUN_DIR}/logs" "${LOCAL_LOG}" /tmp/sudoku_wave_runs
+
+SCP_OPTS="-o IdentitiesOnly=yes -o StrictHostKeyChecking=accept-new"
+[ -f "${HOME}/.ssh/id_ed25519_berkeley" ] && SCP_OPTS="${SCP_OPTS} -i ${HOME}/.ssh/id_ed25519_berkeley"
+
+if ${ENV_LOCAL}/bin/python -u -c "import jax; assert jax.default_backend()=='gpu' or 'cuda' in str(jax.devices()[0]).lower()" 2>/dev/null; then
+  echo "[$(date)] reusing ${ENV_LOCAL}"
+else
+  echo "[$(date)] fetching env tarball"
+  rm -rf "${ENV_LOCAL}"
+  scp ${SCP_OPTS} "gandalf.berkeley.edu:${TARBALL_GANDALF}" "${TARBALL_LOCAL}"
+  tar xzf "${TARBALL_LOCAL}" -C /tmp
+  rm -f "${TARBALL_LOCAL}"
+fi
+
+export PY=${ENV_LOCAL}/bin/python
+export LD_LIBRARY_PATH=\
+${ENV_LOCAL}/lib/python3.9/site-packages/nvidia/cudnn/lib:\
+${ENV_LOCAL}/lib/python3.9/site-packages/nvidia/cublas/lib:\
+${ENV_LOCAL}/lib/python3.9/site-packages/nvidia/cuda_runtime/lib:\
+${ENV_LOCAL}/lib/python3.9/site-packages/nvidia/cuda_nvrtc/lib:\
+${ENV_LOCAL}/lib/python3.9/site-packages/nvidia/nccl/lib:\
+${LD_LIBRARY_PATH:-}
+
+${PY} -u -c "import jax; print(jax.__version__, jax.devices(), jax.default_backend())"
+
+need_inst=0
+for f in train_assignments.npy train_starts.npy train_counts.npy \
+         test_assignments.npy test_starts.npy test_counts.npy; do
+  [ -s "${INST_DIR}/${f}" ] || need_inst=1
+done
+if [ "${need_inst}" = 1 ]; then
+  echo "[$(date)] pulling instances from ${GANDALF_INST}"
+  mkdir -p "${INST_DIR}"
+  rsync -a --progress -e "ssh ${SCP_OPTS}" \
+    "${GANDALF_INST}/" "${INST_DIR}/"
+fi
+for f in train_assignments.npy train_starts.npy train_counts.npy \
+         test_assignments.npy test_starts.npy test_counts.npy; do
+  [ -s "${INST_DIR}/${f}" ] || { echo "missing ${INST_DIR}/${f}" >&2; exit 1; }
+done
+
+need_cand=0
+for f in train_cand_masks.npy test_cand_masks.npy; do
+  [ -s "${CAND_DIR}/${f}" ] || need_cand=1
+done
+if [ "${need_cand}" = 1 ]; then
+  echo "[$(date)] pulling s12 masks from ${GANDALF_CAND}"
+  mkdir -p "${CAND_DIR}"
+  rsync -a -e "ssh ${SCP_OPTS}" "${GANDALF_CAND}/" "${CAND_DIR}/" || true
+fi
+for f in train_cand_masks.npy test_cand_masks.npy; do
+  [ -s "${CAND_DIR}/${f}" ] || { echo "missing ${CAND_DIR}/${f}" >&2; exit 1; }
+done
+
+export SUDOKU_RESUME=0
+export SUDOKU_START_STAGE=1
+export SUDOKU_MAX_STAGE=12
+export SUDOKU_LATENT_SLOTS=12
+export SUDOKU_RECURRENT=1
+export SUDOKU_CAND_SLOT_MODE=depth
+export SUDOKU_PASSES_PER_STAGE=1
+export SUDOKU_AUX_WEIGHT=0.0
+export SUDOKU_LEVEL_BALANCED=0
+export SUDOKU_DATA_CURRICULUM=none
+export SUDOKU_PLATEAU_STEPS=20000
+export SUDOKU_PLATEAU_DELTA=0.005
+export SUDOKU_PATIENCE=80000
+export SUDOKU_MIN_STAGE_STEPS=8000
+export SUDOKU_PROMOTE_ACC=0.85
+export SUDOKU_PROMOTE_LOC=0.70
+export SUDOKU_MAX_STEPS="${SUDOKU_MAX_STEPS:-800000}"
+export SUDOKU_EVAL_EVERY=2000
+export SUDOKU_SAVE_EVERY=10000
+export SUDOKU_CKPT_KEEP=3
+export SUDOKU_LR=0.0002
+export SUDOKU_DROPOUT=0.2
+export SUDOKU_WD=0.005
+export SUDOKU_TRAIN_PATH="${SCRATCH_ROOT}/sudoku-code/datasets/train_sudoku_puzzles.npy"
+export SUDOKU_TEST_PATH="${SCRATCH_ROOT}/sudoku-code/datasets/test_sudoku_puzzles.npy"
+export SUDOKU_TRAIN_CAND="${CAND_DIR}/train_cand_masks.npy"
+export SUDOKU_TEST_CAND="${CAND_DIR}/test_cand_masks.npy"
+export SUDOKU_INSTANCE_DIR="${INST_DIR}"
+export XLA_PYTHON_CLIENT_MEM_FRACTION=0.9
+
+# Adaptive repair: if stage t's in-set rate drops more than 0.03 below
+# the value it graduated at, replay that depth until it recovers.
+export SUDOKU_BACKTRACK=1
+export SUDOKU_BACKTRACK_MODE=adaptive
+export SUDOKU_BACKTRACK_MARGIN=0.03
+export SUDOKU_BACKTRACK_MAX_REPAIR_STEPS=4000
+export SUDOKU_BACKTRACK_MIN_FRONTIER_STEPS=8000
+export SUDOKU_BACKTRACK_MAX_REPAIR_FRACTION=0.25
+export SUDOKU_BACKTRACK_GRAD_DECAY=0.05
+export SUDOKU_BACKTRACK_FRONTIER_MIX=1
+
+cd "${RUN_DIR}"
+(
+  while true; do sleep 900
+    rsync -a "${LOCAL_LOG}.log" "${RUN_DIR}/logs/w12_inst_latent_bt.log" 2>/dev/null || true
+  done
+) &
+SYNC_PID=$!
+
+echo "[$(date)] starting w12_inst_latent_bt from scratch"
+echo "  K=12 recurrent=1 bt=adaptive aux=0 instance_dir=${INST_DIR}"
+CUDA_VISIBLE_DEVICES=0 ${PY} -u -m train.main \
+  --workdir="${LOCAL_LOG}" --exp_name="w12_inst_latent_bt" \
+  > "${LOCAL_LOG}.log" 2>&1
+EC=$?
+kill ${SYNC_PID} 2>/dev/null || true
+rsync -a "${LOCAL_LOG}.log" "${RUN_DIR}/logs/w12_inst_latent_bt.log" 2>/dev/null || true
+echo "[$(date)] w12_inst_latent_bt exit ${EC}"
+exit ${EC}

code/show_instance_stages.py

@@ -0,0 +1,94 @@
+"""Compare what a data instance looks like at different stages, for a few cells.
+
+Tracks the same handful of cells across stages so the shrinking candidate sets
+and the resulting instances line up column by column.
+"""
+import argparse
+
+import numpy as np
+
+import superposition_instances as SI
+
+
+def fmt(ds):
+    return "{" + ",".join(str(d) for d in ds) + "}"
+
+
+def cell_name(c):
+    return f"r{c // 9}c{c % 9}"
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--masks", default="datasets_multicandidate_s12/test_cand_masks.npy")
+    ap.add_argument("--puzzle", type=int, default=2)
+    ap.add_argument("--stages", type=int, nargs="+", default=[1, 5, 11])
+    ap.add_argument("--ncells", type=int, default=6)
+    ap.add_argument("--show-instances", type=int, default=8)
+    ap.add_argument("--max-confine", type=int, default=3)
+    ap.add_argument("--max-instances", type=int, default=32)
+    args = ap.parse_args()
+
+    masks = np.load(args.masks, mmap_mode="r")
+
+    # Pick cells that are still undetermined at the *latest* requested stage
+    # before the final one, so the same cells are interesting at every stage.
+    pick_stage = max(s for s in args.stages if s < masks.shape[1] - 1)
+    m_pick = np.array(masks[args.puzzle, pick_stage]).astype(np.uint16)
+    _, empties = SI.split_cells(m_pick)
+    cells = sorted(empties, key=lambda c: -len(SI.digits_of(m_pick[c])))[:args.ncells]
+    cells = sorted(cells)
+
+    print(f"puzzle {args.puzzle}, tracking {len(cells)} cells: "
+          + ", ".join(cell_name(c) for c in cells))
+
+    for s in args.stages:
+        mask = np.array(masks[args.puzzle, s]).astype(np.uint16)
+        forced, empties = SI.split_cells(mask)
+        res = SI.instances_for_stage(mask, max_confine=args.max_confine,
+                                     max_instances=args.max_instances, seed=7)
+        inst = res["instances"]
+
+        print()
+        print("=" * 88)
+        print(f"STAGE {s}   {len(empties)} cells still undetermined, "
+              f"mean {res['mean_width']:.2f} candidates each")
+        print(f"  generation: {res['n_instances']} instances kept, "
+              f"{res['n_rejected']} ruled out, "
+              f"grid-wide coverage {100*res['coverage']:.1f}%")
+        print("=" * 88)
+
+        print("  candidate set at this stage:")
+        for c in cells:
+            ds = SI.digits_of(mask[c])
+            tag = "  (determined)" if len(ds) == 1 else ""
+            print(f"      {cell_name(c)}  {fmt(ds)}{tag}")
+
+        if not len(inst):
+            print("  no instances")
+            continue
+
+        print(f"\n  first instance, written out as the (row, col, value) triples "
+              f"the model would emit:")
+        print("      " + "  ".join(
+            f"({c//9},{c%9},{inst[0][c]})" for c in cells))
+
+        k = min(args.show_instances, len(inst))
+        print(f"\n  value assigned to each cell, across the first {k} instances:")
+        print("      inst   " + "  ".join(f"{cell_name(c):>6}" for c in cells))
+        for i in range(k):
+            print(f"      {i:>4}   " + "  ".join(
+                f"{inst[i][c]:>6}" for c in cells))
+
+        print(f"\n  coverage over all {res['n_instances']} instances:")
+        for c in cells:
+            ds = set(SI.digits_of(mask[c]))
+            seen = set(int(a[c]) for a in inst)
+            miss = sorted(ds - seen)
+            status = "all candidates seen" if not miss else f"never seen: {fmt(miss)}"
+            print(f"      {cell_name(c)}  candidates {fmt(sorted(ds)):<16} "
+                  f"seen {fmt(sorted(seen)):<16} {status}")
+
+
+if __name__ == "__main__":
+    main()

code/show_superposition_example.py

@@ -0,0 +1,102 @@
+"""Worked example of superposition-instance generation for one puzzle/stage.
+
+Prints the candidate sets, the dependencies read off the mask, the instances that
+survived, the ones that were ruled out, and the resulting per-cell coverage.
+Display is restricted to one unit so the table is readable; generation always
+runs on the full grid.
+"""
+import argparse
+
+import numpy as np
+
+import superposition_instances as SI
+
+NAMES = ([f"row {r}" for r in range(9)] + [f"col {c}" for c in range(9)]
+         + [f"box {b}" for b in range(9)])
+
+
+def fmt(cells):
+    return "{" + ",".join(str(d) for d in cells) + "}"
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--masks", default="datasets_multicandidate_s12/test_cand_masks.npy")
+    ap.add_argument("--puzzle", type=int, default=2)
+    ap.add_argument("--stage", type=int, default=0)
+    ap.add_argument("--unit", type=int, default=14, help="unit index to display (14 = col 5)")
+    ap.add_argument("--max-confine", type=int, default=3)
+    ap.add_argument("--max-instances", type=int, default=16)
+    args = ap.parse_args()
+
+    masks = np.load(args.masks, mmap_mode="r")
+    mask = np.array(masks[args.puzzle, args.stage]).astype(np.uint16)
+
+    unit = SI.UNITS[args.unit]
+    forced, empties = SI.split_cells(mask)
+
+    print("=" * 78)
+    print(f"puzzle {args.puzzle}, stage {args.stage}, displaying {NAMES[args.unit]}")
+    print("=" * 78)
+    print(f"grid: {len(forced)} determined cells, {len(empties)} undetermined, "
+          f"mean candidate count {np.mean([len(SI.digits_of(mask[c])) for c in empties]):.2f}")
+
+    print(f"\ncandidate sets in {NAMES[args.unit]}:")
+    for c in unit:
+        r, cc = divmod(c, 9)
+        ds = SI.digits_of(mask[c])
+        tag = "determined" if len(ds) == 1 else ""
+        print(f"    r{r}c{cc}  {fmt(ds):<22} {tag}")
+
+    all_deps = SI.dependencies_from_mask(mask, max_confine=9)
+    conf_deps = SI.dependencies_from_mask(mask, max_confine=args.max_confine)
+    print(f"\ndependencies over the whole grid:")
+    print(f"    {len(all_deps)} total (every unit, every unplaced digit)")
+    print(f"    {len(conf_deps)} with confinement size <= {args.max_confine} "
+          f"(these are the ones enforced)")
+    sizes = {}
+    for _, S in all_deps:
+        sizes[len(S)] = sizes.get(len(S), 0) + 1
+    print("    by confinement size: "
+          + ", ".join(f"|S|={k}: {v}" for k, v in sorted(sizes.items())))
+
+    print(f"\ndependencies enforced inside {NAMES[args.unit]}:")
+    shown = 0
+    for d, S in conf_deps:
+        if not set(S) <= set(unit):
+            continue
+        locs = " or ".join(f"r{c//9}c{c%9}" for c in S)
+        print(f"    digit {d} must be placed at {locs}")
+        shown += 1
+    if shown == 0:
+        print("    (none at this confinement threshold)")
+
+    res = SI.instances_for_stage(mask, max_confine=args.max_confine,
+                                 max_instances=args.max_instances, seed=0)
+
+    print(f"\ngeneration: {res['n_instances']} instances kept, "
+          f"{res['n_rejected']} ruled out over {res['n_attempts']} attempts")
+    print(f"coverage: {res['n_covered']}/{res['n_pairs']} "
+          f"(cell, candidate) pairs = {100*res['coverage']:.1f}%")
+
+    inst = res["instances"]
+    if len(inst):
+        print(f"\nwhat each instance assigned in {NAMES[args.unit]}:")
+        hdr = "    inst  " + "  ".join(f"r{c//9}c{c%9}" for c in unit)
+        print(hdr)
+        for i, a in enumerate(inst):
+            print(f"    {i:>4}  " + "  ".join(f"{a[c]:>4}" for c in unit))
+
+        print(f"\nper-cell coverage in {NAMES[args.unit]}:")
+        for c in unit:
+            ds = set(SI.digits_of(mask[c]))
+            seen = set(int(a[c]) for a in inst)
+            miss = sorted(ds - seen)
+            r, cc = divmod(c, 9)
+            status = "complete" if not miss else f"missing {fmt(miss)}"
+            print(f"    r{r}c{cc}  candidates {fmt(sorted(ds)):<22} "
+                  f"seen {fmt(sorted(seen)):<22} {status}")
+
+
+if __name__ == "__main__":
+    main()

code/smoke_instance_loader.py

@@ -0,0 +1,74 @@
+"""Sanity-check: instance rewrite keeps clues, varies empties, stays in-set."""
+import os
+import sys
+
+import numpy as np
+
+sys.path.insert(0, os.path.join(os.path.dirname(__file__), "wavecurriculum_run"))
+os.environ.setdefault("CUDA_VISIBLE_DEVICES", "")
+
+from train.data import CurriculumState, SudokuDataset  # noqa: E402
+
+
+class Cfg:
+    seed = 7
+    seq_order = "solver-order"
+    num_latent_slots = 12
+    latent_token_id = 10
+    curriculum_max_stage = 12
+    data_curriculum = "none"
+    cand_slot_mode = "depth"
+    passes_per_stage = 1
+    level_balanced_sampling = 0
+    train_puzzle_path = "datasets/train_sudoku_puzzles.npy"
+    test_puzzle_path = "datasets/test_sudoku_puzzles.npy"
+    train_cand_masks_path = "datasets_multicandidate_s12/train_cand_masks.npy"
+    test_cand_masks_path = "datasets_multicandidate_s12/test_cand_masks.npy"
+    instance_dir = "datasets_superposition"
+    train_meta_path = None
+    test_meta_path = None
+
+
+def main():
+    cur = CurriculumState(stage=1, max_stage=12)
+    ds = SudokuDataset(Cfg(), train=True, curriculum=cur)
+    assert ds.instances is not None
+    assert ds.cand_masks is not None
+
+    n_check = 8
+    for stage in (0, 5, 11):
+        cur.stage = stage + 1
+        clue_ok = clue_tot = 0
+        empty_in = empty_tot = empty_eq = 0
+        for idx in range(n_check):
+            base = ds.train_inputs[idx].copy()
+            si = int(ds.train_start_index[idx, 0])
+            rewritten = ds.apply_instance(base, idx, stage)
+            sol = ds.train_puzzles[idx]
+            mask = np.array(ds.cand_masks[idx, stage])
+            orig = {(int(t[0]), int(t[1])): int(t[2])
+                    for t in base.reshape(81, 3)}
+            for t in rewritten[: 3 * si].reshape(-1, 3):
+                r, c, v = int(t[0]), int(t[1]), int(t[2])
+                clue_tot += 1
+                clue_ok += int(v == orig[(r, c)] == int(sol[r * 9 + c]))
+            for t in rewritten[3 * si:].reshape(-1, 3):
+                r, c, v = int(t[0]), int(t[1]), int(t[2])
+                bits = int(mask[r * 9 + c])
+                empty_tot += 1
+                empty_in += int(1 <= v <= 9 and (bits >> (v - 1)) & 1)
+                empty_eq += int(v == int(sol[r * 9 + c]))
+        print(f"stage {stage}: clues {clue_ok}/{clue_tot}  "
+              f"in-set {empty_in}/{empty_tot}  "
+              f"==sol {empty_eq}/{empty_tot}")
+        assert clue_ok == clue_tot, "clue values must stay the original clues"
+        assert empty_in == empty_tot, "every empty must stay in the stage set"
+        if stage == 11:
+            assert empty_eq == empty_tot, "stage 11 must be the unique solution"
+        else:
+            assert empty_eq < empty_tot, "early stages must still be ambiguous"
+    print("SMOKE OK")
+
+
+if __name__ == "__main__":
+    main()

code/superposition_instances.py

@@ -0,0 +1,274 @@
+"""Build superposition instances from a stage's candidate masks.
+
+An *instance* is one concrete assignment: exactly one digit per cell, each drawn
+from that cell's candidate set at that stage. Across the instances for a stage,
+every (cell, candidate) pair should appear at least once, so the candidate set is
+recoverable from the instances without ever being supervised as a set.
+
+Build chronology (per puzzle, per stage):
+  1. read the candidate sets from the stage mask
+  2. derive the live dependencies from that same mask
+  3. propose an instance
+  4. drop it if it disobeys any dependency
+  5. repeat 3-4 until every (cell, candidate) pair is covered
+
+Dependencies are recomputed from the mask rather than recorded when a technique
+fires. The stage-k mask is the cumulative product of every technique applied up
+to stage k, so a confinement visible in the mask *is* the disjunction those
+techniques established, and it stays visible exactly as long as it is live.
+
+The confinement threshold matters. For a unit U and a digit d not yet placed in
+U, S(U,d) is the set of cells of U that can still hold d, and the dependency is
+"d must be placed somewhere in S(U,d)". Enforcing this for every unit and every
+digit is full Sudoku unit coverage, which forces the unique solution and leaves
+no superposition at all. Restricting to small |S(U,d)| keeps only the genuinely
+locked disjunctions.
+"""
+import numpy as np
+
+POPCOUNT = np.array([bin(i).count("1") for i in range(512)], dtype=np.int8)
+
+
+def build_units():
+    """The 27 units, each a list of 9 cell ids (cell = r*9 + c)."""
+    units = []
+    for r in range(9):
+        units.append([r * 9 + c for c in range(9)])
+    for c in range(9):
+        units.append([r * 9 + c for r in range(9)])
+    for br in range(0, 9, 3):
+        for bc in range(0, 9, 3):
+            units.append([(br + i) * 9 + (bc + j)
+                          for i in range(3) for j in range(3)])
+    return units
+
+
+UNITS = build_units()
+
+
+def digits_of(m):
+    m = int(m)
+    return [d for d in range(1, 10) if m & (1 << (d - 1))]
+
+
+def dependencies_from_mask(mask, max_confine=9):
+    """Live disjunctions readable from one stage's mask.
+
+    Returns a list of (digit, cells) meaning "digit must be placed in one of
+    these cells". Only digits not already pinned in the unit are included, and
+    only when the confinement size is at most max_confine.
+    """
+    deps = []
+    for unit in UNITS:
+        pinned = 0
+        for c in unit:
+            m = int(mask[c])
+            if m and (m & (m - 1)) == 0:
+                pinned |= m
+        for d in range(1, 10):
+            bit = 1 << (d - 1)
+            if pinned & bit:
+                continue
+            S = [c for c in unit if int(mask[c]) & bit]
+            if 1 <= len(S) <= max_confine:
+                deps.append((d, S))
+    return deps
+
+
+def split_cells(mask):
+    """(forced cell -> digit, list of undetermined cells)."""
+    forced, empties = {}, []
+    for c in range(81):
+        m = int(mask[c])
+        if m == 0:
+            continue
+        if m & (m - 1) == 0:
+            forced[c] = m.bit_length()
+        else:
+            empties.append(c)
+    return forced, empties
+
+
+def _violated(assign, deps):
+    return [i for i, (d, S) in enumerate(deps)
+            if not any(assign[c] == d for c in S)]
+
+
+def _violations(assign, deps):
+    return len(_violated(assign, deps))
+
+
+def build_dep_index(deps):
+    """cell -> [(dep index, digit that dep wants), ...]"""
+    cell_deps = [[] for _ in range(81)]
+    for i, (d, S) in enumerate(deps):
+        for c in S:
+            cell_deps[c].append((i, d))
+    return cell_deps
+
+
+def repair(assign, deps, forced, rng, max_iter=200, cell_deps=None):
+    """Min-conflicts repair: repeatedly take a violated disjunction and give its
+    digit to whichever of its cells breaks the fewest other disjunctions.
+
+    Purely a proposal-quality step. The accept/reject test still runs afterwards
+    and is the only thing that decides whether an instance enters the dataset.
+    Violation counts are maintained incrementally, so each move costs only the
+    handful of disjunctions that touch the cell being changed.
+    """
+    if not deps:
+        return True
+    if cell_deps is None:
+        cell_deps = build_dep_index(deps)
+    counts = [sum(1 for c in S if assign[c] == d) for d, S in deps]
+    nviol = sum(1 for x in counts if x == 0)
+
+    def apply(c, new):
+        nonlocal nviol
+        old = assign[c]
+        if old == new:
+            return
+        for i, d in cell_deps[c]:
+            if d == old:
+                counts[i] -= 1
+                if counts[i] == 0:
+                    nviol += 1
+            elif d == new:
+                if counts[i] == 0:
+                    nviol -= 1
+                counts[i] += 1
+        assign[c] = new
+
+    for _ in range(max_iter):
+        if nviol == 0:
+            return True
+        bad = [i for i, x in enumerate(counts) if x == 0]
+        d, S = deps[bad[rng.integers(len(bad))]]
+        free = [c for c in S if c not in forced]
+        if not free:
+            return False
+        best, best_v, prev = None, None, {}
+        for c in free:
+            prev[c] = assign[c]
+            before = nviol
+            apply(c, d)
+            v = nviol
+            apply(c, prev[c])
+            assert nviol == before
+            if best_v is None or v < best_v:
+                best, best_v = c, v
+        apply(best, d)
+    return nviol == 0
+
+
+def propose(mask, forced, empties, deps, uncovered, rng, aware=True):
+    """One candidate instance. `aware` first satisfies the confined
+    disjunctions, then fills the rest preferring not-yet-covered pairs.
+    Without it, every cell is an independent draw from its candidate set."""
+    assign = np.zeros(81, dtype=np.int8)
+    for c, d in forced.items():
+        assign[c] = d
+    taken = set(forced)
+
+    if aware and deps:
+        order = sorted(range(len(deps)), key=lambda i: len(deps[i][1]))
+        for i in order:
+            d, S = deps[i]
+            if any(assign[c] == d for c in S):
+                continue
+            free = [c for c in S if c not in taken]
+            if not free:
+                continue          # unsatisfiable in this proposal; test catches it
+            c = free[rng.integers(len(free))]
+            assign[c] = d
+            taken.add(c)
+
+    for c in empties:
+        if c in taken:
+            continue
+        cands = digits_of(mask[c])
+        unc = [d for d in cands if (c, d) in uncovered]
+        pool = unc if unc else cands
+        assign[c] = pool[rng.integers(len(pool))]
+    return assign
+
+
+def instances_for_stage(mask, max_confine=3, max_instances=64,
+                        max_attempts=400, seed=0, aware=True, max_repair=200,
+                        require_new=True, patience=60):
+    """Run the generate/test/loop for one stage.
+
+    Returns a dict with the instances and the statistics the caller wants:
+    how many were produced, how many were ruled out, and how much of each
+    candidate set the survivors cover.
+    """
+    rng = np.random.default_rng(seed)
+    forced, empties = split_cells(mask)
+    deps = dependencies_from_mask(mask, max_confine)
+
+    all_pairs = {(c, d) for c in empties for d in digits_of(mask[c])}
+    uncovered = set(all_pairs)
+
+    cell_deps = build_dep_index(deps)
+    accepted, rejected, attempts = [], 0, 0
+    seen = set()
+    if not empties:
+        # Fully determined stage (the last one): the single instance is the
+        # solution itself, which the loop below would never emit.
+        only = np.zeros(81, dtype=np.int8)
+        for c, d in forced.items():
+            only[c] = d
+        accepted.append(only)
+    since_new = 0
+    while (uncovered and attempts < max_attempts
+           and len(accepted) < max_instances and since_new < patience):
+        attempts += 1
+        inst = propose(mask, forced, empties, deps, uncovered, rng, aware)
+        if max_repair:
+            repair(inst, deps, forced, rng, max_repair, cell_deps)
+        key = inst.tobytes()
+        if _violations(inst, deps) != 0 or key in seen:
+            rejected += 1
+            since_new += 1
+            continue
+        seen.add(key)
+        gained = [(c, int(inst[c])) for c in empties
+                  if (c, int(inst[c])) in uncovered]
+        if require_new and not gained:
+            # Valid but redundant: adds no candidate the set already lacks.
+            since_new += 1
+            continue
+        since_new = 0
+        accepted.append(inst)
+        for pair in gained:
+            uncovered.discard(pair)
+
+    covered = len(all_pairs) - len(uncovered)
+    widths = [len(digits_of(mask[c])) for c in empties]
+    # A disjunction says "digit d appears in S", never "at most once", so an
+    # instance may place the same digit twice in a unit. Count how often.
+    dups = []
+    for a in accepted:
+        n = 0
+        for unit in UNITS:
+            seen_d = {}
+            for c in unit:
+                seen_d[a[c]] = seen_d.get(a[c], 0) + 1
+            n += sum(v - 1 for v in seen_d.values() if v > 1)
+        dups.append(n)
+    return {
+        "instances": np.array(accepted, dtype=np.int8) if accepted
+                     else np.zeros((0, 81), dtype=np.int8),
+        "n_instances": len(accepted),
+        "n_rejected": rejected,
+        "n_attempts": attempts,
+        "n_deps": len(deps),
+        "dep_sizes": [len(S) for _, S in deps],
+        "n_empty": len(empties),
+        "mean_width": float(np.mean(widths)) if widths else 0.0,
+        "n_pairs": len(all_pairs),
+        "n_covered": covered,
+        "coverage": covered / len(all_pairs) if all_pairs else 1.0,
+        "hit_cap": bool(uncovered),
+        "mean_dups": float(np.mean(dups)) if dups else 0.0,
+    }

code/train/data.py

@@ -0,0 +1,537 @@
+# coding=utf-8
+# Copyright 2024 The Google Research Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Data loading procedure for Othello and Sudoku game.
+"""
+
+import itertools
+import os
+import pickle
+
+import jax
+import numpy as np
+import tensorflow as tf
+from tensorflow.compat.v1 import gfile
+
+import pdb
+
+
+class CurriculumState:
+    """Mutable curriculum stage shared between the train loop and the sampler.
+
+    stage t (1..max_stage) is a REASONING-DEPTH stage: t latent slots are
+    active and the first t wave snapshots are supervised. Every puzzle is
+    available from step 0 -- difficulty is not gated.
+
+    The difficulty tag is deliberately unused as a curriculum axis. It takes
+    only 6 values (3..8 = hardest solver-strategy digit needed by any cell), so
+    it cannot express a 12-step ladder, it is uncorrelated with puzzle size
+    (r=-0.003 vs empty-cell count), and it explains only ~19% of the variance
+    in solver round count. The ladder is over propagation depth instead.
+    """
+
+    def __init__(self, stage=1, max_stage=12):
+        self.stage = stage
+        self.max_stage = max_stage
+
+    def unlocked_levels(self):
+        """All levels, always. Kept for logging/compat with per-level reports."""
+        return list(range(3, 9))
+
+
+def compute_puzzle_levels(strategy_codes, start_index):
+    """Per-puzzle difficulty level = max strategy digit over solution cells.
+
+    strategy_codes: (N, 81) int64 chain codes (digits = strategy applications).
+    Clue cells have code 0. Level clipped to [3, 8] (rare all-lone-single
+    puzzles fold into level 3).
+    """
+    n = len(strategy_codes)
+    levels = np.zeros(n, dtype=np.int32)
+    chunk = 200000
+    for lo in range(0, n, chunk):
+        x = strategy_codes[lo:lo + chunk].astype(np.int64).copy()
+        m = np.zeros_like(x)
+        while x.any():
+            np.maximum(m, x % 10, out=m)
+            x //= 10
+        levels[lo:lo + chunk] = m.max(axis=1).astype(np.int32)
+    return np.clip(levels, 3, 8)
+
+
+def create_dataset(config, bs, train, curriculum=None):
+    """Create Sudoku dataset according to the config.
+
+    Args:
+        config: a config object containing the hyparameters for the dataset
+            creation.
+        bs: batch size
+        train: whether the dataset is for train or eval
+
+    Returns:
+        a tf.data.Dataset object
+    """
+    ds, output_types, output_shapes = None, None, None
+    ds = SudokuDataset(config, train=train, curriculum=curriculum)
+    # Each example is (sequence with latent slots, solution, start_index,
+    # difficulty level, per-slot candidate-set targets, round-bin). The
+    # round-bin is 0 unless the round-count data curriculum is enabled.
+    K = int(getattr(config, "num_latent_slots", 0))
+    output_types = (tf.int32, tf.int32, tf.int32, tf.int32, tf.int32, tf.int32)
+    output_shapes = (
+        tf.TensorShape([config.seq_len]),
+        tf.TensorShape([config.block_size]),
+        tf.TensorShape([1]),
+        tf.TensorShape([1]),
+        tf.TensorShape([K, config.block_size]),
+        tf.TensorShape([1]),
+    )
+
+    # Create a tf.data.Dataset object from the generator.
+    tf_ds = tf.data.Dataset.from_generator(
+        generator=ds, output_types=output_types, output_shapes=output_shapes)
+
+    # Repeat the dataset indefinitely.
+    tf_ds = tf_ds.repeat()
+    # Shuffle the dataset with a buffer size of 8 * bs and a seed of 0.
+    tf_ds = tf_ds.shuffle(8 * config.minibatch_size, seed=0)
+    # Batch the dataset with a batch size of bs.
+    tf_ds = tf_ds.batch(bs)
+    return tf_ds
+
+
+
+def prepare_tf_data(xs):
+    """Convert a input batch from tf Tensors to numpy arrays."""
+    def _prepare(x):
+        return x._numpy()  # pylint: disable=protected-access
+    
+    return jax.tree_map(_prepare, xs)
+
+
+def create_iter(config, bs, train, curriculum=None):
+    tf_ds = create_dataset(config, bs, train=train, curriculum=curriculum)
+    it = map(prepare_tf_data, tf_ds)
+    return it
+
+class SudokuDataset:
+    """Sudoku dataset."""
+    def __init__(self, config, train=True, curriculum=None):
+        self.config = config
+        self.train = train
+        self.curriculum = curriculum
+        self.num_latent_slots = int(getattr(config, "num_latent_slots", 0))
+        self.latent_token_id = int(getattr(config, "latent_token_id", 10))
+        self.rng = np.random.RandomState(config.seed if hasattr(config, "seed") else 0)
+        self.preprocess_sudoku()
+        self._load_candidate_masks()
+        self._load_round_bins()
+        self._load_instances()
+
+    def _load_instances(self):
+        """Load superposition instances: one concrete assignment per row.
+
+        Replaces the multi-hot candidate target with ordinary value tokens. For
+        a puzzle at stage s there are several assignments, each picking one digit
+        per cell from that cell's stage-s candidate set, so the candidate set is
+        recoverable across instances instead of being supervised as a set.
+
+        assignments: (M, 81) uint8, cell = r*9+c
+        starts/counts: (N, S) row range for each (puzzle, stage)
+        """
+        self.instances = None
+        d = getattr(self.config, "instance_dir", None)
+        if not d:
+            return
+        split = "train" if self.train else "test"
+        self.instances = np.load(
+            os.path.join(d, f"{split}_assignments.npy"), mmap_mode="r")
+        self.inst_starts = np.load(os.path.join(d, f"{split}_starts.npy"))
+        self.inst_counts = np.load(os.path.join(d, f"{split}_counts.npy"))
+        self.inst_stages = int(self.inst_starts.shape[1])
+        print(f"[inst] loaded {split} instances {self.instances.shape} "
+              f"over {self.inst_starts.shape[0]} puzzles, "
+              f"{self.inst_stages} stages", flush=True)
+
+    def instance_stage(self):
+        """Stage whose instances this example should target.
+
+        Bound to the curriculum stage so the target's ambiguity matches the
+        latent depth: stage t runs t latent slots and supervises the stage-(t-1)
+        assignments, ending at the unique solution when t == S.
+        """
+        S = self.inst_stages
+        if self.curriculum is None:
+            return S - 1
+        t = int(np.clip(self.curriculum.stage, 1, S))
+        return t - 1
+
+    def instance_values(self, idx, stage):
+        """One assignment for (puzzle idx, stage), sampled uniformly."""
+        n = int(self.inst_counts[idx, stage])
+        if n <= 0:
+            return None
+        row = int(self.inst_starts[idx, stage]) + self.rng.randint(n)
+        return np.asarray(self.instances[row])
+
+    def apply_instance(self, seq, idx, stage):
+        """Rewrite the value token of every triple to this instance's digit.
+
+        The (row, col) order is untouched, so the clue block and the solver-order
+        output sequence are exactly as before; only the values change.
+        """
+        vals = self.instance_values(idx, stage)
+        if vals is None:
+            return seq
+        seq = seq.copy()
+        cells = seq[0::3].astype(np.int64) * 9 + seq[1::3].astype(np.int64)
+        seq[2::3] = vals[cells].astype(seq.dtype)
+        return seq
+
+    def _load_round_bins(self):
+        """Load per-puzzle solver round counts and bin them into max_stage
+        equal-count bins, for the round-count DATA curriculum.
+
+        The round count (waves needed to reach the unique solution, 5..38) is
+        the same propagation-depth axis the latent curriculum supervises, but
+        used to order the *puzzles* instead of the supervision. Bin edges are
+        always computed on the train split and reused for eval so a bin index
+        means the same thing in both.
+        """
+        self.round_bins = None
+        self.num_bins = int(getattr(self.config, "curriculum_max_stage", 12))
+        if str(getattr(self.config, "data_curriculum", "none")) != "rounds":
+            return
+        tr_path = getattr(self.config, "train_meta_path", None)
+        path = tr_path if self.train else getattr(
+            self.config, "test_meta_path", None)
+        if not (path and tr_path):
+            raise ValueError(
+                "data_curriculum='rounds' needs SUDOKU_TRAIN_META and "
+                "SUDOKU_TEST_META (the *_meta.npy written by "
+                "staged_candidate_gen.py; column 2 is num_rounds)")
+        rounds = np.load(path, mmap_mode="r")[:, 2].astype(np.int32)
+        # Cut points from the TRAIN split, so a bin index means the same thing
+        # in eval. Round counts are integers with a peaked distribution (mean
+        # 22, sd 4), so raw quantiles collide -- the 12-bin quantiles repeat 20
+        # twice on the full corpus, which would leave a bin permanently empty
+        # and strand its stage with no frontier to measure. Force the cuts
+        # strictly increasing so every bin is reachable.
+        train_rounds = np.load(tr_path, mmap_mode="r")[:, 2].astype(np.int32)
+        edges = np.quantile(train_rounds, np.linspace(0, 1, self.num_bins + 1))
+        cuts = np.round(edges[1:-1]).astype(np.int64)
+        for i in range(1, len(cuts)):
+            if cuts[i] <= cuts[i - 1]:
+                cuts[i] = cuts[i - 1] + 1
+        # bin j (1-based) = stage that first unlocks the puzzle.
+        self.round_bins = np.clip(
+            np.searchsorted(cuts, rounds, side="right") + 1,
+            1, self.num_bins).astype(np.int32)
+        self.bin_index = {b: np.where(self.round_bins == b)[0]
+                          for b in range(1, self.num_bins + 1)}
+        counts = {b: int(len(v)) for b, v in self.bin_index.items()}
+        print(f"[rounds] {'train' if self.train else 'eval'} bin counts:",
+              counts, flush=True)
+        print(f"[rounds] cuts: {cuts.tolist()}  (rounds "
+              f"{int(rounds.min())}..{int(rounds.max())})", flush=True)
+        empty = [b for b, c in counts.items() if c == 0]
+        if empty and self.train:
+            raise ValueError(
+                f"round-bin curriculum has empty train bins {empty}; those "
+                f"stages would have no puzzles and no frontier signal")
+
+    def _load_candidate_masks(self):
+        """Load staged candidate-set masks (N, S, 81) uint16, aligned by puzzle
+        index with the loaded .npy. Row i here == puzzle i in the base file."""
+        if self.train:
+            path = getattr(self.config, "train_cand_masks_path", None)
+        else:
+            path = getattr(self.config, "test_cand_masks_path", None)
+        self.cand_masks = None
+        self.num_stages = 0
+        if path:
+            self.cand_masks = np.load(path, mmap_mode="r")
+            self.num_stages = int(self.cand_masks.shape[1])
+            print(f"[cand] loaded {path} shape {self.cand_masks.shape}", flush=True)
+
+    def slot_budget(self, level):
+        """Number of latent slots this example activates (see cand_slot_mode).
+
+        "depth" mode must agree with the recurrence depth used by the train
+        step, since build_latent_state only ever writes slots [0, num_passes):
+        supervising a slot the recurrence never filled would train the head off
+        an all-zero latent.
+        """
+        K = self.num_latent_slots
+        if getattr(self.config, "cand_slot_mode", "level") == "depth":
+            stage = self.curriculum.stage if self.curriculum is not None \
+                else getattr(self.config, "curriculum_max_stage", 6)
+            pps = int(getattr(self.config, "passes_per_stage", 1))
+            return int(np.clip(pps * stage, 1, K))
+        return int(np.clip(level - 2, 1, K))
+
+    def _slot_stage_targets(self, idx, level, clue_cells=None):
+        """Return (K, 81) int32 candidate bitmasks, one per latent slot.
+
+        The S stored stages are mapped onto the example's k active slots; see
+        cand_slot_mode for the two mappings ("level" re-paces the whole shrink
+        sequence into k slots, "depth" assigns slot j to stage j). Inactive slots
+        (j>=k) default to the final stage; they are masked out of the loss.
+
+        Although the array is laid out over all 81 cell positions (for a fixed
+        batch shape), the supervised targets are only the *empty* cells: clue
+        cells are zeroed out here as a sentinel (a genuine empty cell always has
+        >=1 candidate at every stage), and the loss ignores zero rows. So the
+        effective target per puzzle is (#empty cells) x 9, in solver order."""
+        K = self.num_latent_slots
+        if self.cand_masks is None or K == 0:
+            return np.zeros((K, 81), dtype=np.int32)
+        S = self.num_stages
+        k = self.slot_budget(level)
+        depth_mode = getattr(self.config, "cand_slot_mode", "level") == "depth"
+        stages = self.cand_masks[idx].astype(np.int32)   # (S, 81)
+        out = np.zeros((K, 81), dtype=np.int32)
+        for j in range(K):
+            if j >= k:
+                # Inactive slot: masked out of the loss, value is irrelevant.
+                s = S - 1
+            elif depth_mode:
+                # Identity: slot j holds propagation block j, so growing the
+                # recurrence depth extends the chain instead of re-pacing it.
+                # The solution is only reached at full depth, which is what
+                # makes this a curriculum over reasoning depth.
+                s = min(j, S - 1)
+            else:
+                # Active slots span the full shrink sequence: slot 0 -> stage 0
+                # (widest candidate set, genuinely multi-valued), last active
+                # slot -> final stage (solution). For k==1 the single slot maps
+                # to the WIDEST set (stage 0), not the solution, so even level-3
+                # puzzles give the candidate head a real multi-candidate target
+                # (the LM head still produces the unique answer).
+                s = int(round(j * (S - 1) / max(k - 1, 1)))
+            out[j] = stages[s]
+        # Sentinel-zero the clue cells so only the empty cells are supervised.
+        if clue_cells is not None and len(clue_cells) > 0:
+            out[:, clue_cells] = 0
+        return out
+
+    def _build_level_index(self, levels):
+        """Map difficulty level -> array of puzzle indices."""
+        return {lvl: np.where(levels == lvl)[0] for lvl in range(3, 9)}
+
+    def insert_latent_slots(self, seq, start_index):
+        """Insert K latent placeholder tokens between clues and solution.
+
+        seq: (243,) triple sequence. Returns (243 + K,) sequence:
+        [clues (3*si)] [K placeholders] [solution triples].
+        """
+        k = self.num_latent_slots
+        if k == 0:
+            return seq
+        si3 = 3 * int(start_index)
+        return np.concatenate([
+            seq[:si3],
+            np.full(k, self.latent_token_id, dtype=seq.dtype),
+            seq[si3:],
+        ])
+
+
+    def convert_to_fixed_or_random_order(self, inputs, start_index):
+        """Convert the sequence of moves to either a fixed or random order.
+
+        Args:
+            inputs: a numpy array of shape (num_puzzles, seq_len) containing the
+                sequence of moves for each puzzle
+            start_index: a numpy array of shape (num_puzzles, 1) containing the starting
+                index for each puzzle
+
+        Returns:
+            transformed_input: a numpy array of shape (num_puzzles, seq_len) containing the
+                sequence of moves for each puzzle in either a fixed or random order
+        """
+        transformed_input = np.zeros_like(inputs)
+        
+        for i in range(len(inputs)):
+            cur_seq = inputs[i]
+            cur_start_index = start_index[i, 0]
+            
+            # Split the sequence into input and output prompts
+            inp_prompt = cur_seq[ :(3 * cur_start_index) ].reshape(-1, 3)
+            out_prompt = cur_seq[ (3 * cur_start_index): ].reshape(-1, 3)
+            
+            # Sort the input prompts in a fixed order
+            if self.config.seq_order == "fixed":
+                transformed_input[i, :(3 * cur_start_index) ] = inp_prompt[ np.lexsort( inp_prompt[:, ::-1].T ) ].flatten()
+            # Randomly shuffle the input prompts
+            elif self.config.seq_order == "random":
+                transformed_input[i, :(3 * cur_start_index) ] = np.random.permutation(inp_prompt).flatten()
+            
+            # Sort the output prompts in a fixed order
+            if self.config.seq_order == "fixed":
+                transformed_input[i, (3 * cur_start_index): ] = out_prompt[ np.lexsort( out_prompt[:, ::-1].T ) ].flatten()
+            # Randomly shuffle the output prompts
+            elif self.config.seq_order == "random":
+                transformed_input[i, (3 * cur_start_index): ] = np.random.permutation(out_prompt).flatten()
+        
+        return transformed_input
+
+    def get_puzzles_start_index(self, path):
+        """Get the puzzles, start index, inputs and difficulty levels.
+
+        Returns:
+            inputs: (num_puzzles, 243) move sequences (strategy column removed)
+            puzzles: (num_puzzles, 81) solutions
+            start_index: (num_puzzles, 1) number of clue cells
+            levels: (num_puzzles,) puzzle difficulty level in [3, 8]
+                (= hardest solver-strategy digit needed by any cell)
+        """
+        with gfile.Open(path, "rb") as f:
+            inputs_with_start_index = np.load(f)
+        start_index = inputs_with_start_index[:, 0]  # Get the start index
+
+        rest = inputs_with_start_index[:, 1:]
+        # Strategy chain codes (4th entry of each cell quadruple); keep them to
+        # derive the curriculum difficulty level, then remove from the inputs.
+        strategy_codes = rest.reshape(len(rest), 81, 4)[:, :, 3]
+        levels = compute_puzzle_levels(strategy_codes, start_index)
+        inputs = np.delete( rest, np.arange(81) * 4 + 3, axis=1)
+        
+        puzzles = np.zeros((len(inputs), 81), dtype=np.int8)  # Initialize puzzles
+        for j in range(81):
+            cell_id = inputs[:, 3 * j] * 9 + inputs[:, 3 * j + 1]  # Get the cell id
+            puzzles[np.arange(len(inputs)), cell_id] = inputs[:, 3 * j + 2]  # Set the puzzle
+        
+        return inputs, puzzles, start_index.reshape(-1, 1), levels
+    
+    
+    def preprocess_sudoku(self):
+        """Preprocess the sudoku for train and test datasets.
+        
+        Depending on the `train` flag, this method loads and processes the
+        sudoku puzzles and their start indices from the appropriate paths, and
+        optionally converts them to a fixed or random order based on the 
+        configuration.
+        """
+        if self.train is True:
+            # Load train puzzles, inputs, and start indices
+            (self.train_inputs, self.train_puzzles, self.train_start_index,
+             self.train_levels) = (
+                self.get_puzzles_start_index(self.config.train_puzzle_path)
+            )
+            # Convert train inputs to fixed or random order if specified
+            if self.config.seq_order in {"fixed", "random"}:
+                self.train_inputs = self.convert_to_fixed_or_random_order(self.train_inputs, self.train_start_index)
+            self.level_index = self._build_level_index(self.train_levels)
+            print("train level counts:",
+                  {l: len(v) for l, v in self.level_index.items()}, flush=True)
+        
+        elif self.train is False:
+            # Load evaluation puzzles, inputs, and start indices
+            (self.eval_inputs, self.eval_puzzles, self.eval_start_index,
+             self.eval_levels) = (
+                self.get_puzzles_start_index(self.config.test_puzzle_path)
+            )
+            # Convert evaluation inputs to fixed or random order if specified
+            if self.config.seq_order in {"fixed", "random"}: 
+                self.eval_inputs = self.convert_to_fixed_or_random_order(self.eval_inputs, self.eval_start_index)
+            self.level_index = self._build_level_index(self.eval_levels)
+
+    def __len__(self):
+        if self.train is True:
+            return len(self.train_puzzles)
+        elif self.train is False:
+            return len(self.eval_puzzles)
+        
+    def __getitem__(self, idx):
+        """Returns one example: (sequence with latent slots, solution,
+        start_index, difficulty level).
+
+        The base sequence is 243 tokens of (row, column, value) triples; K
+        latent placeholder tokens are inserted after the clue block, giving
+        243 + K tokens. start_index is the number of clue cells; level in
+        [3, 8] is the hardest solver strategy needed by any cell.
+        """
+        if self.train is True:
+            inputs, puzzles = self.train_inputs, self.train_puzzles
+            start_index, levels = self.train_start_index, self.train_levels
+        else:
+            inputs, puzzles = self.eval_inputs, self.eval_puzzles
+            start_index, levels = self.eval_start_index, self.eval_levels
+
+        base = inputs[idx, :]
+        if self.instances is not None and self.train:
+            # Same input prompt, different output prompt: the clue triples are
+            # untouched (their instance digit is the clue) while the empty cells
+            # take one draw from the stage's candidate sets. Averaged over the
+            # instances the target IS the candidate set, so the superposition is
+            # learned from ordinary next-token CE instead of a set head. Eval
+            # keeps the unique solution: the sequence it scores is generated, and
+            # `puzzles` must stay the ground truth the accuracy is measured on.
+            base = self.apply_instance(base, idx, self.instance_stage())
+        seq = self.insert_latent_slots(base, start_index[idx, 0])
+        # Clue cells = the first `start_index` (r,c,v) triples; their cell ids
+        # are excluded from candidate supervision (only empty cells are scored).
+        si = int(start_index[idx, 0])
+        clue_triples = inputs[idx, :3 * si].reshape(-1, 3)
+        clue_cells = (clue_triples[:, 0] * 9 + clue_triples[:, 1]).astype(np.int64)
+        cand_targets = self._slot_stage_targets(idx, int(levels[idx]), clue_cells)
+        rbin = (int(self.round_bins[idx]) if self.round_bins is not None else 0)
+        return (
+            seq,
+            puzzles[idx, :],
+            start_index[idx],
+            np.array([levels[idx]], dtype=np.int32),
+            cand_targets,
+            np.array([rbin], dtype=np.int32),
+        )
+
+    def _sample_level(self):
+        """Uniform over levels that have puzzles. Used for level-balanced mode
+        and for eval, where per-level reporting needs every level represented."""
+        avail = [l for l in range(3, 9) if len(self.level_index[l]) > 0]
+        return avail[self.rng.randint(len(avail))]
+
+    def _sample_round_gated(self):
+        """Uniform over puzzles whose round-bin is already unlocked (bin<=stage).
+        Reading the stage at yield time lets the pool grow on promotion."""
+        stage = self.curriculum.stage if self.curriculum is not None \
+            else self.num_bins
+        stage = int(np.clip(stage, 1, self.num_bins))
+        pool = np.concatenate([self.bin_index[b] for b in range(1, stage + 1)])
+        return int(pool[self.rng.randint(len(pool))])
+
+    def __call__(self):
+        # Infinite generator. Train draws puzzles uniformly from the whole
+        # corpus (natural difficulty mix, ~68% level 3): the difficulty tag
+        # selects nothing, since the curriculum axis is propagation depth.
+        # Eval stays level-balanced so per-level accuracy is measurable and
+        # comparable across runs.
+        #
+        # The exception is the round-count DATA curriculum, where the train pool
+        # is restricted to puzzles needing at most stage-many propagation waves.
+        # Eval is never gated: it must score the whole corpus at every stage.
+        round_gated = (self.train and self.round_bins is not None)
+        level_balanced = (not self.train) or bool(
+            int(getattr(self.config, "level_balanced_sampling", 0)))
+        n = len(self.train_puzzles) if self.train else len(self.eval_puzzles)
+        while True:
+            if round_gated:
+                idx = self._sample_round_gated()
+            elif level_balanced:
+                idx_arr = self.level_index[self._sample_level()]
+                idx = int(idx_arr[self.rng.randint(len(idx_arr))])
+            else:
+                idx = int(self.rng.randint(n))
+            yield self.__getitem__(idx)

code/train/evaluater.py

@@ -0,0 +1,455 @@
+"""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

code/train/main.py

@@ -0,0 +1,320 @@
+# coding=utf-8
+# Copyright 2024 The Google Research Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Main file for Sudoku GPT experiments."""
+
+import os
+import sys
+
+from absl import app
+from absl import flags
+from absl import logging
+
+from clu import platform
+
+import jax
+import tensorflow as tf
+import wandb
+
+import ml_collections
+from ml_collections import config_flags
+
+from train import train_and_evaluate
+from train import train_backtrack
+
+import pdb
+
+
+
+sys.dont_write_bytecode = True
+
+logging.set_verbosity(logging.INFO)
+
+FLAGS = flags.FLAGS
+
+_WORKDIR = flags.DEFINE_string(
+    'workdir',
+    None,
+    'Directory to store model data.')
+_EXP_NAME = flags.DEFINE_string(
+    'exp_name',
+    None,
+    'Experiment name.')
+_CKPT_LOC = flags.DEFINE_string(
+    'ckpt_loc',
+    None,
+    'Directory to restore model.')
+
+config_flags.DEFINE_config_file(
+    'config',
+    None,
+    'File path to the training hyperparameter configuration.',
+    lock_config=True)
+flags.mark_flags_as_required(['workdir', 'exp_name'])
+
+
+def get_config():
+    """Get the default hyperparameter configuration.
+
+    Returns:
+    A ConfigDict object.
+    """
+  
+    # Common configuration for all experiments.
+    config = ml_collections.ConfigDict()
+
+    # Dataset choice
+    config.dataset = 'sudoku'
+
+    # Sequence order
+    config.seq_order = "solver-order"           ## Choices = ["fixed", "solver-order", "random"]
+
+    # Training related parameters
+    config.max_steps = int(os.environ.get("SUDOKU_MAX_STEPS", 100000))
+    config.dtype = jax.numpy.bfloat16
+    config.minibatch_size = int(os.environ.get("SUDOKU_MINIBATCH", 64))
+
+    # Continuous latent thoughts (ATC / Coconut style) + curriculum.
+    # Set SUDOKU_LATENT_SLOTS=0 for the no-latent control baseline (plain
+    # transformer + difficulty curriculum, no recurrence, no candidate head).
+    config.num_latent_slots = int(os.environ.get("SUDOKU_LATENT_SLOTS", 6))
+    config.latent_token_id = 10      # placeholder id for inactive slots
+    # Recurrent latent feedback. 1 = full ATC/Coconut recurrence (slot k built
+    # from slot k-1's hidden). 0 = NO recurrence: the K stage slots + shared
+    # candidate head + curriculum + BCE supervision are all kept, but the latent
+    # hidden state is never fed back, so each stage is an independent parallel
+    # readout with no carried state (the "stagewise, no-recurrence" control).
+    config.recurrent_latent = int(os.environ.get("SUDOKU_RECURRENT", 1))
+    # Backtracking (stage-replay) training: when 1, use the train_backtrack loop
+    # which interleaves frontier-stage batches with replays of earlier stages at
+    # their MATCHED recurrence depth (num_passes=t for stage t) to mitigate
+    # forgetting. backtrack_prob = fraction of steps that are replay steps.
+    config.backtrack = int(os.environ.get("SUDOKU_BACKTRACK", 0))
+    # Replay strategy: "prob" = fixed-probability replay of a random earlier stage
+    # (backtrack_prob). "adaptive" = deficit-driven controller that only replays a
+    # stage once its held-out accuracy has regressed below its graduation value by
+    # backtrack_margin, repairs the most-deficient stage first with a FULL
+    # frontier-style step, and returns once it recovers or hits the step cap.
+    config.backtrack_mode = os.environ.get("SUDOKU_BACKTRACK_MODE", "prob")
+    config.backtrack_prob = float(os.environ.get("SUDOKU_BACKTRACK_PROB", 0.5))
+    config.backtrack_margin = float(os.environ.get("SUDOKU_BACKTRACK_MARGIN", 0.03))
+    config.backtrack_max_repair_steps = int(
+        os.environ.get("SUDOKU_BACKTRACK_MAX_REPAIR_STEPS", 6000))
+    # After leaving a repair episode, require this many frontier steps before
+    # another repair can trigger (prevents thrashing / starves-the-frontier).
+    config.backtrack_min_frontier_steps = int(
+        os.environ.get("SUDOKU_BACKTRACK_MIN_FRONTIER_STEPS", 0))
+    # While the frontier stage's own level accuracy is below this target, use a
+    # stricter (larger) repair trigger so stage-6 training is prioritized until
+    # it is itself "good". 0 disables.
+    config.backtrack_frontier_target_acc = float(
+        os.environ.get("SUDOKU_BACKTRACK_FRONTIER_TARGET_ACC", 0.0))
+    # Optional explicit graduation refs for inherited stages when resuming at
+    # stage > 1. Comma-separated floats for stages 1..start-1, e.g.
+    # "0.462,0.487,0.534,0.597,0.636". Empty = seed from first eval (old behavior).
+    config.backtrack_grad_acc_seed = os.environ.get("SUDOKU_GRAD_ACC_SEED", "")
+    # Max fraction of steps (since resume) that may be spent in repair mode.
+    # 0 disables the cap. Typical fair setting: 0.25.
+    config.backtrack_max_repair_fraction = float(
+        os.environ.get("SUDOKU_BACKTRACK_MAX_REPAIR_FRACTION", 0.0))
+    # Soft graduation refs: each eval, if current < grad, decay
+    #   grad <- (1-d)*grad + d*current
+    # so chronic mild regression is forgiven; sharp drops still trigger.
+    # 0 disables (hard bars). Typical: 0.05–0.1.
+    config.backtrack_grad_decay = float(
+        os.environ.get("SUDOKU_BACKTRACK_GRAD_DECAY", 0.0))
+    # After this absolute step, force frontier-only (no new repairs).
+    # 0 disables. Use for "BT early, then freeze".
+    config.backtrack_freeze_after_step = int(
+        os.environ.get("SUDOKU_BACKTRACK_FREEZE_AFTER", 0))
+    # Frontier steps sample uniformly over ALL unlocked levels (standard-loop
+    # data mix) instead of only the frontier level. 1 = fair "standard+repairs"
+    # recipe (default); 0 = legacy frontier-level-only behavior.
+    config.backtrack_frontier_mix = int(
+        os.environ.get("SUDOKU_BACKTRACK_FRONTIER_MIX", 1))
+    config.curriculum_start_stage = int(os.environ.get("SUDOKU_START_STAGE", 1))
+    # Stage s = reasoning DEPTH s: s latent slots active, wave snapshots 1..s
+    # supervised. Difficulty is NOT gated -- every puzzle is available from
+    # step 0. Cap max_stage (e.g. =start_stage) to hold a single depth.
+    config.curriculum_max_stage = int(os.environ.get("SUDOKU_MAX_STAGE", 12))
+    # Promotion is gated on the candidate-set accuracy of the DEEPEST active
+    # slot, restricted to cells that changed from the previous snapshot.
+    config.promote_acc_threshold = float(os.environ.get("SUDOKU_PROMOTE_ACC", 0.85))
+    # Instance arm: inset is only scored on location-correct cells. Without a
+    # location floor, a 0.90 inset on 2% loc_acc would promote from a handful
+    # of cells. Stage-0 chance inset is ~3.7/9 ≈ 0.41, so 0.85 is the set bar
+    # and this is the "model actually emits the solver-order cells" bar.
+    config.promote_loc_threshold = float(os.environ.get("SUDOKU_PROMOTE_LOC", 0.70))
+    config.promote_patience_steps = int(os.environ.get("SUDOKU_PATIENCE", 8000))
+    config.min_stage_steps = int(os.environ.get("SUDOKU_MIN_STAGE_STEPS", 2000))
+    # Plateau promotion: advance once the frontier depth stops improving, rather
+    # than on a fixed timer. A stage is "done" when it has not gained
+    # plateau_delta over its best accuracy for plateau_steps steps. This is the
+    # primary rule; the accuracy threshold is a fast path for mastery and
+    # promote_patience_steps is a hard cap so a stuck stage cannot stall
+    # training forever. Set plateau_steps=0 to disable.
+    config.plateau_steps = int(os.environ.get("SUDOKU_PLATEAU_STEPS", 20000))
+    config.plateau_delta = float(os.environ.get("SUDOKU_PLATEAU_DELTA", 0.005))
+    # Train-time difficulty balancing. 0 (default) = draw puzzles uniformly from
+    # the corpus, so the difficulty tag selects nothing. 1 = uniform over the 6
+    # levels, which upsamples level 8 from 1.8% to 16.7% of batches. Eval is
+    # always level-balanced so per-level accuracy stays measurable.
+    config.level_balanced_sampling = int(
+        os.environ.get("SUDOKU_LEVEL_BALANCED", 0))
+    # Data curriculum over the puzzle POOL (as opposed to the latent-depth
+    # curriculum over the supervision).
+    #   "none"   = every puzzle available from step 0.
+    #   "rounds" = stage t admits only puzzles whose solver round count falls in
+    #              the first t of max_stage equal-count bins (rounds span 5..38,
+    #              so 12 bins give a genuinely smooth 12-step ladder). This is
+    #              the same propagation-depth axis the latent arm supervises,
+    #              which makes the two arms directly comparable. Needs
+    #              SUDOKU_TRAIN_META / SUDOKU_TEST_META.
+    config.data_curriculum = os.environ.get("SUDOKU_DATA_CURRICULUM", "none")
+    config.train_meta_path = os.environ.get("SUDOKU_TRAIN_META", "") or None
+    config.test_meta_path = os.environ.get("SUDOKU_TEST_META", "") or None
+
+    # Model related parameters
+    config.block_size = 81
+    config.seq_len = 3 * config.block_size + config.num_latent_slots
+    config.vocab_size = 11
+
+    # Model architecture
+    config.num_heads = 8
+    config.num_layers = 8
+    config.emb_dim = 576
+    config.qkv_dim = 576
+    config.mlp_dim = 6 * config.emb_dim
+    config.dropout_rate = float(os.environ.get("SUDOKU_DROPOUT", 0.2))
+    config.attention_dropout_rate = float(
+        os.environ.get("SUDOKU_ATTN_DROPOUT",
+                       os.environ.get("SUDOKU_DROPOUT", 0.2)))
+
+    # Training hyperparameters
+    config.learning_rate = float(os.environ.get("SUDOKU_LR", 0.0002))  # Base learning rate.
+    config.end_lr_factor = float(os.environ.get("SUDOKU_END_LR_FACTOR", 0.2))
+    config.warmup_tokens = int(os.environ.get("SUDOKU_WARMUP", 10000))
+    config.weight_decay = float(os.environ.get("SUDOKU_WD", 0.005))
+    # Resume from a checkpoint (set SUDOKU_RESUME=1 and pass --ckpt_loc=<path>).
+    config.resume_training = os.environ.get("SUDOKU_RESUME", "0") == "1"
+
+    # Other hyperparameters
+    config.seed = 7
+    config.save_checkpoint = os.environ.get("SUDOKU_SAVE_CKPT", "1") == "1"
+    config.save_every_steps = int(os.environ.get("SUDOKU_SAVE_EVERY", 10000))
+    # How many checkpoints to retain. Large default so per-stage checkpoints are
+    # never rolled off (disk is plentiful; ~0.5GB each).
+    config.ckpt_keep = int(os.environ.get("SUDOKU_CKPT_KEEP", 100))
+    config.use_wandb = False
+    config.wandb_project_name = 'sudoku'
+
+    # Evaluation related parameters
+    config.eval_every_steps = int(os.environ.get("SUDOKU_EVAL_EVERY", 2000))
+    config.eval_epochs = int(os.environ.get("SUDOKU_EVAL_EPOCHS", 5))
+
+    # Path to dataset
+    config.train_puzzle_path = os.environ.get(
+        "SUDOKU_TRAIN_PATH", "datasets/train_sudoku_puzzles.npy")
+    config.train_candidate_path = "datasets/train_sudoku_puzzles_candidate.npy"
+    config.test_puzzle_path = os.environ.get(
+        "SUDOKU_TEST_PATH", "datasets/test_sudoku_puzzles.npy")
+    config.test_candidate_path = "datasets/test_sudoku_puzzles_candidate.npy"
+
+    # Staged multi-candidate supervision (per-latent-slot BCE targets).
+    # Empty string disables cand-mask loading (useful for K=0 baselines).
+    config.train_cand_masks_path = os.environ.get(
+        "SUDOKU_TRAIN_CAND", "datasets_multicandidate/train_cand_masks.npy") or None
+    config.test_cand_masks_path = os.environ.get(
+        "SUDOKU_TEST_CAND", "datasets_multicandidate/test_cand_masks.npy") or None
+    # Superposition-instance targets. When set, the output prompt's value tokens
+    # come from one sampled stage-k assignment instead of the unique solution:
+    # the input prompt (clues) is fixed and the same puzzle recurs with different
+    # legal completions, so the candidate set is represented across the batch
+    # rather than supervised as a multi-hot set. Setting this should go with
+    # SUDOKU_AUX_WEIGHT=0 (candidate head off) -- the masks are then read only
+    # for the in-set metric. Empty string = classic single-solution targets.
+    config.instance_dir = os.environ.get("SUDOKU_INSTANCE_DIR", "") or None
+    # Weight of the auxiliary candidate-set BCE loss relative to the LM CE loss.
+    config.aux_cand_weight = float(os.environ.get("SUDOKU_AUX_WEIGHT", 1.0))
+    # Positive-class weight inside the candidate BCE (counters the sparsity of
+    # the multi-hot masks so the head doesn't collapse to predicting all-zeros).
+    config.aux_pos_weight = float(os.environ.get("SUDOKU_CAND_POS_WEIGHT", 5.0))
+
+    # How many latent slots each example activates.
+    #   "level" = k = clip(level-2, 1, K).  Difficulty-matched, but puzzle level
+    #             explains only ~19% of the variance in solver round count, so
+    #             most examples leave the majority of the K slots inert (a
+    #             level-3 puzzle activates ONE slot for ~21 rounds of work).
+    #   "depth"  = k = num_passes, uniform over the batch. Every slot the
+    #             recurrence actually fills is active and supervised, and the
+    #             curriculum advances reasoning depth rather than puzzle level.
+    config.cand_slot_mode = os.environ.get("SUDOKU_CAND_SLOT_MODE", "depth")
+    # Latent passes granted per curriculum stage: num_passes = min(pps*stage, K).
+    # pps=2 with K=12 reaches all 12 slots by stage 6.
+    config.passes_per_stage = int(os.environ.get("SUDOKU_PASSES_PER_STAGE", 1))
+    # Loss weight for candidate cells that did NOT change from the previous
+    # stage. Consecutive stages are highly redundant (at 12 stages ~97% of the
+    # target bits are copies of the previous slot), so plain BCE is dominated by
+    # echoing the previous slot. <1.0 down-weights the copied cells and puts the
+    # gradient on the digits actually eliminated at this stage. 1.0 = off.
+    config.aux_delta_bg = float(os.environ.get("SUDOKU_CAND_DELTA_BG", 1.0))
+
+    return config
+
+
+def main(argv):
+    if len(argv) > 1:
+        raise app.UsageError('Too many command-line arguments.')
+
+    # # Hide any GPUs from TensorFlow. Otherwise TF might reserve memory and make
+    # # it unavailable to JAX.
+    tf.config.experimental.set_visible_devices([], 'GPU')
+
+    cfgs = get_config()
+    if cfgs.resume_training: 
+        assert _CKPT_LOC.value is not None
+  
+    if cfgs.use_wandb:
+        wandb.init(project=cfgs.wandb_project_name, name=_EXP_NAME.value, config=cfgs)
+
+    logging.info('JAX process: %d / %d', jax.process_index(), jax.process_count())
+    logging.info('JAX local devices: %r', jax.local_devices())
+
+    # Add a note so that we can tell which task is which JAX host.
+    # (Depending on the platform task 0 is not guaranteed to be host 0)
+    platform.work_unit().set_task_status(f'process_index: {jax.process_index()}, '
+                                       f'process_count: {jax.process_count()}')
+    platform.work_unit().create_artifact(platform.ArtifactType.DIRECTORY,
+                                       _WORKDIR.value, 'workdir')
+  
+    logging.info(cfgs)
+
+    cfgs.workdir = _WORKDIR.value
+    cfgs.ckpt_loc = _CKPT_LOC.value
+    if int(getattr(cfgs, "backtrack", 0)):
+        if str(getattr(cfgs, "backtrack_mode", "prob")) == "adaptive":
+            train_backtrack.train_and_evaluate_backtrack_adaptive(cfgs, _WORKDIR.value)
+        else:
+            train_backtrack.train_and_evaluate_backtrack(cfgs, _WORKDIR.value)
+    else:
+        train_and_evaluate.train_and_evaluate(cfgs, _WORKDIR.value)
+
+    if cfgs.use_wandb:
+        wandb.finish()
+
+
+if __name__ == '__main__':
+    jax.config.config_with_absl()
+    app.run(main)

code/train/model.py

@@ -0,0 +1,232 @@
+# coding=utf-8
+# Copyright 2024 The Google Research Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Model Architecture."""
+
+import functools
+from typing import Any, Callable
+
+from flax import linen as nn
+from flax import struct
+from jax import numpy as jnp
+
+@struct.dataclass
+class TransformerConfig:
+    """Global hyperparameters used to minimize obnoxious kwarg plumbing."""
+    vocab_size: int = 1
+    dtype: Any = jnp.float32
+    emb_dim: int = 512
+    num_heads: int = 8
+    num_layers: int = 6
+    qkv_dim: int = 512
+    mlp_dim: int = 2048
+    seq_len: int = 2048  # Maximum sequence length
+    dropout_rate: float = 0.1
+    attention_dropout_rate: float = 0.1
+    deterministic: bool = False
+    num_latent_slots: int = 0  # K continuous latent thought slots
+    # When False, the fed-back latent thought vectors are NOT injected into the
+    # slot positions (they keep their placeholder token embedding). This turns
+    # the K slots into static, parallel per-stage readouts with NO carried
+    # recurrent state -- the "stagewise supervision, no recurrence" control.
+    inject_latents: bool = True
+
+    
+class TransformerBlock(nn.Module):
+    config: Any = None
+
+    def setup(self):
+        self.vocab_size = self.config.vocab_size
+        self.emb_dim = self.config.emb_dim
+        self.num_layers = self.config.num_layers
+
+    @nn.compact
+    def __call__(self, inputs, causal_mask_inputs, training=True):
+        """
+        Transformer Block call function.
+
+        Args:
+            inputs: Input tensor.
+            causal_mask_inputs: Causal mask for the inputs.
+            training: Whether the model is in training mode.
+
+        Returns:
+            Transformed tensor after self-attention and MLP layers.
+        """
+        
+        x = inputs + nn.SelfAttention(
+            num_heads=self.config.num_heads, dtype=self.config.dtype,
+            qkv_features=self.config.qkv_dim,
+            kernel_init=nn.initializers.xavier_uniform(),
+            bias_init=nn.initializers.normal(stddev=1e-6), 
+            use_bias=False, broadcast_dropout=False,
+            dropout_rate=self.config.attention_dropout_rate, normalize_qk=True,
+            deterministic=self.config.deterministic)(inputs, causal_mask_inputs)
+
+        def mlp(x):
+            """
+            Multi-Layer Perceptron function.
+
+            Args:
+                x: Input tensor.
+
+            Returns:
+                Transformed tensor after applying MLP layers.
+            """
+            dense_with_init = functools.partial(
+                nn.Dense,
+                kernel_init=nn.initializers.xavier_uniform(),
+                bias_init=nn.initializers.normal(stddev=1e-6)
+                )
+            x = dense_with_init(features=self.config.mlp_dim)(x)
+            x = nn.gelu(x)
+            x = dense_with_init(features=self.config.emb_dim)(x)
+            x = nn.Dropout(rate=self.config.dropout_rate,
+                        deterministic=self.config.deterministic)(x)
+            return x
+
+        x = x + mlp(x)
+        return x
+
+
+class TransformerLMHeadModel(nn.Module):
+    config: Any = None
+
+    def setup(self):
+        self.vocab_size = self.config.vocab_size
+        self.emb_dim = self.config.emb_dim
+        self.num_layers = self.config.num_layers
+
+    @nn.compact
+    def __call__(self, inputs, latent_values=None, latent_positions=None,
+                 latent_active=None, training=True):
+        """
+        Transformer LM Head call function.
+
+        Args:
+            inputs: Input token ids (batch, seq).
+            latent_values: Optional (batch, K, emb_dim) continuous thought
+                vectors (raw last-layer hiddens fed back, Coconut/ATC style).
+            latent_positions: Optional (batch, K) int positions of the latent
+                slots in the sequence (per-example, after the clue block).
+            latent_active: Optional (batch, K) bool; slot j uses the projected
+                latent vector when True, otherwise keeps the placeholder
+                token embedding.
+            training: Whether the model is in training mode.
+
+        Returns:
+            (logits, hidden): LM logits and final (post-LayerNorm) hidden
+            states, the latter used to build the next continuous thought.
+        """
+        batch_size, seq_size = inputs.shape
+
+        causal_mask_x = nn.make_causal_mask(inputs, dtype=self.config.dtype)
+
+        # Embed the input tensor using a learnable embedding matrix.
+        embed_with_init = functools.partial(
+            nn.Embed, embedding_init=nn.initializers.normal(stddev=0.02))
+        token_embeddings = embed_with_init(
+            num_embeddings=self.config.vocab_size,
+            features=self.config.emb_dim,
+        )(inputs)
+
+        # Check the shape of the embedded tensor.
+        assert token_embeddings.shape == (batch_size, seq_size,
+                                      self.config.emb_dim)
+
+        # Continuous latent thoughts: project fed-back hidden states and
+        # scatter them into the latent slot positions, replacing the
+        # placeholder token embedding (position embeddings still added below).
+        # Skipped when inject_latents is False (no-recurrence control: the slots
+        # stay as static placeholders and the candidate heads become parallel
+        # per-stage readouts with no carried state).
+        if latent_values is not None and self.config.inject_latents:
+            proj = nn.Dense(features=self.config.emb_dim,
+                            kernel_init=nn.initializers.xavier_uniform(),
+                            name="latent_proj_in")(latent_values)
+            proj = nn.gelu(proj)
+            proj = nn.Dense(features=self.config.emb_dim,
+                            kernel_init=nn.initializers.xavier_uniform(),
+                            name="latent_proj_out")(proj)
+            bidx = jnp.arange(batch_size)[:, None]
+            cur = token_embeddings[bidx, latent_positions]
+            new = jnp.where(latent_active[..., None],
+                            proj.astype(cur.dtype), cur)
+            token_embeddings = token_embeddings.at[
+                bidx, latent_positions].set(new)
+
+        # Initialize the positional embedding variable.
+        pos_embedding_variable = self.variable(
+            "params",
+            "position_embeddings",
+            jnp.zeros,
+            (self.config.seq_len, self.config.emb_dim),
+        )
+
+        # Slice the positional embedding array to the correct sequence length.
+        pos_embeddings = pos_embedding_variable.value[:seq_size, :]
+
+        # Check the shape of the positional embedding array.
+        output_tuple = (pos_embeddings.shape, token_embeddings.shape[1:])
+        assert pos_embeddings.shape == token_embeddings.shape[1:], output_tuple
+
+        # Add the positional embeddings to the token embeddings.
+        x = token_embeddings + pos_embeddings[None, :, :]
+
+        # Apply dropout to the input.
+        x = nn.Dropout(rate=self.config.dropout_rate,
+                    deterministic=self.config.deterministic)(x)
+
+        # Apply the Transformer layers. remat (gradient checkpointing) keeps
+        # the multi-pass latent recurrence within GPU memory under full BPTT.
+        RematBlock = nn.remat(TransformerBlock)
+        for i in range(self.num_layers):
+            x = RematBlock(config=self.config)(
+                    x, causal_mask_x, training=training)
+      
+            self.sow('intermediates', 'feature_' + str(i), x)
+
+        # Apply the final layer normalization.
+        x = nn.LayerNorm()(x)
+
+        # Apply the LM head.
+        logits = nn.Dense(features=self.config.vocab_size,
+                        kernel_init=nn.initializers.xavier_uniform(),
+                        bias_init=nn.initializers.normal(stddev=1e-6),
+                        use_bias=False)(x)
+
+        # Check the shape of the output tensor.
+        assert logits.shape == (batch_size, seq_size, self.config.vocab_size)
+
+        # ---- Auxiliary multi-candidate head ----
+        # Each latent slot reads its (post-LayerNorm) hidden and predicts the
+        # full 81x9 candidate grid for its reasoning stage. Trained with BCE
+        # (independent per-digit sigmoids = candidate-set membership), NOT
+        # softmax, so multiple digits can be "on" at intermediate stages.
+        cand_logits = None
+        if latent_positions is not None:
+            bidx = jnp.arange(batch_size)[:, None]
+            slot_hidden = x[bidx, latent_positions]          # (bs, K, emb)
+            h = nn.Dense(features=self.config.emb_dim,
+                         kernel_init=nn.initializers.xavier_uniform(),
+                         name="cand_head_in")(slot_hidden)
+            h = nn.gelu(h)
+            cand_logits = nn.Dense(features=81 * 9,
+                                   kernel_init=nn.initializers.xavier_uniform(),
+                                   name="cand_head_out")(h)   # (bs, K, 729)
+            cand_logits = cand_logits.reshape(
+                batch_size, -1, 81, 9)                        # (bs, K, 81, 9)
+
+        return logits, x, cand_logits

code/train/train_and_evaluate.py

@@ -0,0 +1,386 @@
+# coding=utf-8
+# Copyright 2024 The Google Research Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""This file contains function that coordinates the training and evaluation of the model."""
+
+import functools
+import math
+import os
+
+from absl import logging
+from clu import metric_writers
+from flax import jax_utils
+from flax import linen as nn
+from flax.training import checkpoints
+import jax
+from jax import random
+import jax.numpy as jnp
+import numpy as np
+import tensorflow as tf
+import wandb
+
+from train import data
+from train import evaluater
+from train import model
+from train import trainer
+
+
+
+def log_hyperparams_tb(
+    config, model_config, initial_variables, tf_summary_writer
+):
+    """Log hyperparameters to TensorBoard.
+
+    Args:
+        config: experiment's ConfigDict
+        model_config: model's ConfigDict
+        initial_variables: initial hyperparameter values
+        tf_summary_writer: SummaryWriter object.
+
+    Returns:
+        The SummaryWriter object and the config.
+    """
+    # Calculate the total number of model parameters
+    config.num_model_parameters = sum(
+        x.size for x in jax.tree_util.tree_leaves(initial_variables)
+    )
+
+    # Convert hyperparameters to tensors
+    config_hyperparameters = [
+        tf.convert_to_tensor([k, str(v)]) for k, v in config.items()
+    ]
+    model_config_hyperparameters = [
+        tf.convert_to_tensor([k, str(v)])
+        for k, v in model_config.__dict__.items()
+    ]
+
+    # Log model hyperparameters to TensorBoard
+    with tf_summary_writer.as_default():
+        tf.summary.text(
+            "Model hyperparameters", tf.stack(model_config_hyperparameters), step=0
+        )
+        tf.summary.text(
+            "Config hyperparameters", tf.stack(config_hyperparameters), step=0
+        )
+
+    return tf_summary_writer, config
+
+
+
+def train_and_evaluate(config, workdir):
+    """The training and evaluation loops for the model.
+
+    Args:
+        config: experiment's config dictionary.
+        workdir: directory to use for logging.
+    """
+    # Orbax checkpointing requires an absolute path.
+    workdir = os.path.abspath(workdir)
+
+    logging.info("Creating training and evaluator dataset iterator")
+    curriculum = data.CurriculumState(
+        stage=int(getattr(config, "curriculum_start_stage", 1)),
+        max_stage=int(getattr(config, "curriculum_max_stage", 6)))
+    train_data_iter = data.create_iter(
+        config, config.minibatch_size, train=True, curriculum=curriculum)
+    eval_data_iter = data.create_iter(config, config.minibatch_size, train=False)
+
+    logging.info("Finished creating training dataset iterator")
+
+    model_config = model.TransformerConfig(
+        dtype=config.dtype,
+        vocab_size=config.vocab_size,
+        seq_len=config.seq_len,
+        num_heads=config.num_heads,
+        num_layers=config.num_layers,
+        emb_dim=config.emb_dim,
+        qkv_dim=config.qkv_dim,
+        mlp_dim=config.mlp_dim,
+        dropout_rate=config.dropout_rate,
+        attention_dropout_rate=config.attention_dropout_rate,
+        deterministic=False,
+        num_latent_slots=int(config.num_latent_slots),
+        inject_latents=bool(int(getattr(config, "recurrent_latent", 1))),
+    )
+
+    logging.info("train_config: %s", str(model_config.__dict__))
+    print(str(model_config.__dict__), flush=True)
+
+    rng = jax.random.PRNGKey(config.seed)
+    rng, init_rng, inference_rng = random.split(rng, num=3)
+
+    # Initialize the model and get initial variables. Dummy latent arguments
+    # are provided so the latent projector parameters are created at init.
+    rng, dropout_rng = jax.random.split(rng)
+    input_shape = (config.minibatch_size, config.seq_len)
+    net = model.TransformerLMHeadModel(model_config)
+    rng_keys = {"params": init_rng, "dropout": dropout_rng}
+    K = int(config.num_latent_slots)
+    dummy_latents = jnp.zeros(
+        (config.minibatch_size, K, config.emb_dim), model_config.dtype)
+    dummy_positions = jnp.zeros((config.minibatch_size, K), jnp.int32)
+    dummy_active = jnp.zeros((config.minibatch_size, K), bool)
+    sample_out, initial_variables = jax.jit(
+        net.init_with_output
+        )(rng_keys, jnp.ones(input_shape, jnp.int32), dummy_latents,
+          dummy_positions, dummy_active)
+    
+    state, lr_scheduler_fn = trainer.get_state(config, net, initial_variables)
+    # Resume-and-extend support: when resuming, start the training loop at the
+    # restored optimizer step (not 0) so a larger config.max_steps continues the
+    # cosine LR schedule cleanly instead of re-running steps or overshooting.
+    start_step = 0
+    if config.resume_training:
+        state = checkpoints.restore_checkpoint(config.ckpt_loc, state)
+        start_step = int(state.step)
+        print("----------Restored model from", config.ckpt_loc,
+              f"at step {start_step}-----------")
+
+    writer = metric_writers.create_default_writer(
+        workdir, asynchronous=False, just_logging=(jax.process_index() > 0))
+    tf_summary_writer = tf.summary.create_file_writer(workdir)
+
+    logging.info("config: %s", str(config.__dict__))
+    state = jax_utils.replicate(state)
+
+    dropout_rngs = jax.random.split(rng, jax.local_device_count())
+
+    def make_p_train_step(num_passes):
+        return jax.pmap(
+            functools.partial(
+                trainer.train_step,
+                config=model_config,
+                hyperparams=config,
+                learning_rate_fn=lr_scheduler_fn,
+                num_passes=num_passes),
+            axis_name="batch",
+            donate_argnums=(0,))
+
+    # num_passes = latent recurrence depth for this stage, capped by the number
+    # of latent slots (K=0 -> 0 passes -> no-latent control baseline).
+    # When recurrence is disabled (control), force 0 passes so no hidden state
+    # is ever fed back: the slots become static, independent per-stage readouts.
+    recurrent = bool(int(getattr(config, "recurrent_latent", 1)))
+    passes_per_stage = int(getattr(config, "passes_per_stage", 1))
+    def passes_for(stage):
+        return min(passes_per_stage * stage, K) if recurrent else 0
+    p_train_step = make_p_train_step(passes_for(curriculum.stage))
+
+    p_eval_step = jax.pmap(functools.partial(evaluater.eval_step,
+                                             config=model_config.replace(deterministic=True)),
+                                             axis_name="batch")
+    
+    hooks, report_progress, train_metrics = trainer.get_metrics_report_progress(
+      config, workdir, writer)
+
+    tf_summary_writer, config = log_hyperparams_tb(
+        config, model_config, initial_variables, tf_summary_writer
+    )
+
+    promote_threshold = float(getattr(config, "promote_acc_threshold", 0.85))
+    promote_loc_threshold = float(getattr(config, "promote_loc_threshold", 0.70))
+    promote_patience = int(getattr(config, "promote_patience_steps", 8000))
+    min_stage_steps = int(getattr(config, "min_stage_steps", 2000))
+    plateau_steps = int(getattr(config, "plateau_steps", 0))
+    plateau_delta = float(getattr(config, "plateau_delta", 0.005))
+    instance_mode = bool(getattr(config, "instance_dir", None))
+    per_stage_inset = {}
+    # Best frontier-depth accuracy seen in the current stage, and the step it
+    # was last improved: the plateau detector's state.
+    stage_best_acc = -1.0
+    stage_best_step = start_step
+    ckpt_keep = int(getattr(config, "ckpt_keep", 100))
+    # Protected directory for per-stage checkpoints (never rolled off).
+    stage_ckpt_dir = os.path.join(workdir, "stage_ckpts")
+    stage_started_at = start_step
+
+    with metric_writers.ensure_flushes(writer):
+        for step in range(start_step, config.max_steps):
+            if step%10000 == 0:
+                print("Step:", step, flush=True)
+
+            state, metrics = trainer.train_one_step(p_train_step, config, state, 
+                                                    step, dropout_rngs, train_data_iter)
+            
+            for h in hooks:
+                h(step)
+
+            if math.isnan(metrics["loss"][0]):
+                print("The loss function became nan: This might be due to the choice of hyperparameters.")
+                break
+
+            if step % config.eval_every_steps == 0:
+                eval_metrics = evaluater.get_eval_metrics(
+                    state, eval_data_iter, p_eval_step, config)
+                per_level = eval_metrics.pop("per_level_acc")
+                per_slot = eval_metrics.pop("per_slot_acc", {})
+                per_slot_ch = eval_metrics.pop("per_slot_acc_changed", {})
+                per_stage_inset = eval_metrics.pop("per_stage_inset_acc", {})
+                per_bin = eval_metrics.pop("per_bin_acc", {})
+                def _m(key):
+                    v = eval_metrics.get(key, [])
+                    return round(float(np.mean(v)), 4) if len(v) else -1.0
+                print(step,
+                      "stage", curriculum.stage,
+                      "loss", round(float(metrics["loss"].mean()), 4),
+                      "ce", round(float(metrics["ce_loss"].mean()), 4),
+                      "aux_bce", round(float(metrics["aux_loss"].mean()), 4),
+                      "| val_acc", _m("acc"),
+                      "loc_acc", _m("loc_acc"),
+                      "val|loc", _m("val_given_loc_acc"),
+                      "| cand_bit_acc", _m("cand_bit_acc"),
+                      "cand_set_acc", _m("cand_set_acc"),
+                      "cand_set_chg", _m("cand_set_acc_changed"),
+                      "| inset", round(
+                          float(per_stage_inset.get(curriculum.stage, -1.0)), 4),
+                      "promote_need inset>={:.2f} loc>={:.2f}".format(
+                          promote_threshold, promote_loc_threshold),
+                      "| depth_acc", round(
+                          float(per_slot_ch.get(curriculum.stage, -1.0)), 4),
+                      flush=True)
+                with tf_summary_writer.as_default():
+                    tf.summary.scalar("loss", metrics["loss"].mean(), step=step)
+                    tf.summary.scalar("ce_loss", metrics["ce_loss"].mean(), step=step)
+                    tf.summary.scalar("aux_bce_loss", metrics["aux_loss"].mean(), step=step)
+                    tf.summary.scalar(
+                        "learning rate", metrics["learning_rate"].mean(), step=step
+                        )
+                    tf.summary.scalar("curriculum_stage", curriculum.stage, step=step)
+
+                    log_dict = {'loss': metrics["loss"].mean(), 'learning rate': metrics["learning_rate"].mean()}
+
+                    for key in eval_metrics.keys():
+                        vals = eval_metrics[key]
+                        if not vals:
+                            continue
+                        tf.summary.scalar(
+                            "eval_" + key, np.array(vals).mean(), step=step
+                        )
+                        log_dict[ "eval_" + key ] = np.array(vals).mean()
+
+                    for lvl, v in per_level.items():
+                        if v >= 0:
+                            tf.summary.scalar(f"eval_acc_level_{lvl}", v, step=step)
+                            log_dict[f"eval_acc_level_{lvl}"] = v
+
+                    for s, v in per_slot.items():
+                        if v >= 0:
+                            tf.summary.scalar(f"eval_cand_depth_{s}", v, step=step)
+                            log_dict[f"eval_cand_depth_{s}"] = v
+                    for s, v in per_slot_ch.items():
+                        if v >= 0:
+                            tf.summary.scalar(f"eval_cand_depth_chg_{s}", v, step=step)
+
+                    if config.use_wandb: wandb.log(log_dict, step=step)
+
+                # ---- Curriculum promotion (ATC-style: threshold + patience) ----
+                # Two curricula, one frontier rule. With latent slots the
+                # frontier is the DEEPEST active slot, i.e. the newest wave
+                # snapshot this stage introduced, scored on changed cells only
+                # (unchanged cells are copies of slot j-1 and stay correct for a
+                # head that learned nothing new). With the round-count data
+                # curriculum the frontier is the newest round-BIN, i.e. the
+                # longest-chain puzzles this stage admitted.
+                rounds_curric = str(
+                    getattr(config, "data_curriculum", "none")) == "rounds"
+                has_frontier = int(config.num_latent_slots) > 0 or rounds_curric
+                if has_frontier and curriculum.stage < curriculum.max_stage:
+                    t = curriculum.stage
+                    loc_now = _m("loc_acc")
+                    if instance_mode:
+                        # Inset = "emitted digit is in the stage-t set", scored
+                        # only on location-correct cells. Require BOTH: the
+                        # solver-order locations, and values inside the set.
+                        # Otherwise a high inset on ~2% loc_acc promotes from
+                        # noise. Plateau is also blocked until location is up,
+                        # so we cannot skip stage 1 the way the old patience
+                        # promotions did at 0.32-0.46.
+                        inset_now = per_stage_inset.get(t, -1.0)
+                        frontier_acc = float(inset_now) if inset_now >= 0 else -1.0
+                    elif int(config.num_latent_slots) > 0:
+                        frontier_acc = per_slot_ch.get(t, -1.0)
+                        if frontier_acc < 0:
+                            frontier_acc = per_slot.get(t, -1.0)
+                    else:
+                        frontier_acc = per_bin.get(t, -1.0)
+                        if frontier_acc < 0:
+                            # The frontier bin holds as little as 6.5% of the
+                            # corpus, so a single eval can miss it. Fall back to
+                            # the pooled accuracy over everything unlocked.
+                            seen = [v for b, v in per_bin.items()
+                                    if b <= t and v >= 0]
+                            frontier_acc = (float(np.mean(seen)) if seen
+                                            else -1.0)
+                    steps_in_stage = step - stage_started_at
+                    # A negative accuracy means "not measured this eval". Never
+                    # let that count as evidence: it must not reset the plateau
+                    # tracker, and it must not satisfy the threshold or plateau
+                    # rule. Only the hard patience cap can fire without a
+                    # measurement.
+                    measured = frontier_acc >= 0
+                    if measured and frontier_acc > stage_best_acc + plateau_delta:
+                        stage_best_acc = frontier_acc
+                        stage_best_step = step
+                    loc_ready = (not instance_mode) or loc_now >= promote_loc_threshold
+                    hit_threshold = (measured and loc_ready
+                                     and frontier_acc >= promote_threshold)
+                    stalled = (measured and loc_ready and plateau_steps > 0
+                               and (step - stage_best_step) >= plateau_steps)
+                    patience_over = steps_in_stage >= promote_patience
+                    if steps_in_stage >= min_stage_steps and (
+                            hit_threshold or stalled or patience_over):
+                        reason = ("threshold" if hit_threshold
+                                  else "plateau" if stalled else "patience")
+                        curriculum.stage += 1
+                        stage_started_at = step
+                        stage_best_acc = -1.0
+                        stage_best_step = step
+                        p_train_step = make_p_train_step(passes_for(curriculum.stage))
+                        what = ("round-bin" if int(config.num_latent_slots) == 0
+                                else "depth")
+                        print(f"[curriculum] step {step}: promote to stage "
+                              f"{curriculum.stage} ({reason}; graduated {what} "
+                              f"{t} acc={frontier_acc:.3f} after "
+                              f"{steps_in_stage} steps); "
+                              f"latent passes={curriculum.stage}, "
+                              f"pool/snapshots now 1..{curriculum.stage}",
+                              flush=True)
+                        if config.save_checkpoint:
+                            unrep_state = jax_utils.unreplicate(state)
+                            # Rolling checkpoint in the main workdir.
+                            checkpoints.save_checkpoint_multiprocess(
+                                workdir, unrep_state, step,
+                                keep=ckpt_keep, overwrite=True)
+                            # Protected copy: the model *entering* this stage,
+                            # kept permanently under stage_ckpts/ (never rolled
+                            # off), so every stage's checkpoint survives.
+                            checkpoints.save_checkpoint_multiprocess(
+                                stage_ckpt_dir, unrep_state, step,
+                                keep=100, overwrite=True,
+                                prefix=f"stage{curriculum.stage}_")
+
+            if config.save_checkpoint and step > 0 and step % config.save_every_steps == 0:
+                checkpoints.save_checkpoint_multiprocess(
+                    workdir, jax_utils.unreplicate(state), step,
+                    keep=ckpt_keep, overwrite=True
+                )
+
+        # Final checkpoint at the end of training
+        if config.save_checkpoint:
+            checkpoints.save_checkpoint_multiprocess(
+                workdir, jax_utils.unreplicate(state), config.max_steps,
+                keep=ckpt_keep, overwrite=True)
+
+  

code/train/train_backtrack.py

@@ -0,0 +1,773 @@
+"""Backtracking (stage-replay) training loop for the recurrent latent model.
+
+Motivation
+----------
+Standard curriculum training runs EVERY batch at `num_passes = curriculum.stage`
+(the current frontier depth). When the curriculum advances to stage i, the
+earlier, shallower readouts (stage 1..i-1, each of which should be produced by
+applying the tied recurrent step exactly t times) drift / are forgotten.
+
+Backtracking = interleaved replay with the recurrence depth MATCHED to the
+replayed stage:
+
+    stage t  <->  difficulty level t+2  <->  num_passes = t  <->  t active slots
+
+With cand_slot_mode="depth" and passes_per_stage=pps the depth becomes pps*t and
+the active-slot count follows num_passes rather than the puzzle level. Replay is
+still safe because that mode maps slot j -> stage j: slots [0, pps*t) hold the
+right targets no matter which frontier stage the sampler built the batch for.
+
+  * frontier step : level (i+2) puzzles, num_passes = i   -> trains stage-i readout
+  * backtrack step: level (t+2) puzzles, num_passes = t   -> re-derives stage-t
+                    for a randomly chosen earlier t in {1..i-1}
+
+Each optimisation step is frontier with prob (1 - backtrack_prob) and a backtrack
+step with prob backtrack_prob (t uniform over the earlier stages). Because the
+recurrent step is weight-tied, replaying "apply f exactly t times = stage t"
+keeps f reusable at every depth instead of specialising to the frontier depth.
+
+Everything else (model, data, per-slot BCE + LM CE loss, eval, promotion) is
+reused unchanged from the standard pipeline.
+"""
+
+import functools
+import math
+import os
+
+from absl import logging
+from clu import metric_writers
+from flax import jax_utils
+from flax.training import checkpoints
+from flax.training import common_utils
+import jax
+import numpy as np
+import tensorflow as tf
+
+from train import data
+from train import evaluater
+from train import model
+from train import trainer
+
+
+class StageBatchSampler:
+    """Draws minibatches of a *specific* difficulty level from the train set.
+
+    Reuses SudokuDataset (which loads the puzzles + staged candidate masks and
+    builds a level->indices map) and assembles numpy batches in the exact tuple
+    layout the train step expects: (seq, puzzle, start_index, level, cand_targets).
+    """
+
+    def __init__(self, config, seed=0):
+        self.ds = data.SudokuDataset(config, train=True)
+        self.rng = np.random.RandomState(seed)
+        # Only levels that actually have puzzles.
+        self.available = {lvl: idx for lvl, idx in self.ds.level_index.items()
+                          if len(idx) > 0}
+
+    def has_level(self, level):
+        return level in self.available
+
+    def _gather(self, idxs):
+        seqs, puzzles, starts, levels, cands = [], [], [], [], []
+        for idx in idxs:
+            seq, puzzle, start_index, lvl, cand = self.ds.__getitem__(int(idx))
+            seqs.append(seq)
+            puzzles.append(puzzle)
+            starts.append(start_index)
+            levels.append(lvl)
+            cands.append(cand)
+        return (
+            np.stack(seqs).astype(np.int32),
+            np.stack(puzzles).astype(np.int32),
+            np.stack(starts).astype(np.int32),
+            np.stack(levels).astype(np.int32),
+            np.stack(cands).astype(np.int32),
+        )
+
+    def sample_all(self, bs):
+        """Batch drawn uniformly from the whole corpus, matching the standard
+        loop's sampler. Difficulty is not a curriculum axis, so every step --
+        frontier or repair -- uses this same distribution and differs only in
+        the recurrence depth it trains at."""
+        return self._gather(self.rng.randint(len(self.ds), size=bs))
+
+    def sample(self, level, bs):
+        idx_pool = self.available[level]
+        idxs = idx_pool[self.rng.randint(len(idx_pool), size=bs)]
+        return self._gather(idxs)
+
+    def sample_mixed(self, levels, bs):
+        """Batch drawn uniformly over `levels` (level first, then a puzzle of
+        that level) — the exact distribution the standard curriculum sampler
+        uses. Frontier steps must use this, NOT single-level batches: training
+        exclusively on the frontier level starves every easier level and is a
+        different (worse) recipe than the standard loop, not "standard + BT"."""
+        levels = [l for l in levels if l in self.available]
+        idxs = []
+        for _ in range(bs):
+            lvl = levels[self.rng.randint(len(levels))]
+            pool = self.available[lvl]
+            idxs.append(int(pool[self.rng.randint(len(pool))]))
+        return self._gather(idxs)
+
+
+def _run_train_step(step_fn, state, batch_tuple, dropout_rngs):
+    """Shard a single numpy batch and run one (pmapped) train step."""
+    inputs, _, start_index, levels, cand_targets = batch_tuple
+    inputs = common_utils.shard(jax.tree_util.tree_map(np.asarray, inputs))
+    start_index = common_utils.shard(jax.tree_util.tree_map(np.asarray, start_index))
+    levels = common_utils.shard(jax.tree_util.tree_map(np.asarray, levels))
+    cand_targets = np.asarray(cand_targets)
+    _nd = jax.local_device_count()
+    cand_targets = cand_targets.reshape(
+        (_nd, cand_targets.shape[0] // _nd) + cand_targets.shape[1:])
+    state, metrics, _ = step_fn(
+        state, inputs, start_index, levels, cand_targets, dropout_rng=dropout_rngs)
+    return state, metrics
+
+
+def _prepare_backtrack(config, workdir):
+    """Build the model, state, pmapped steps, sampler, and writers shared by the
+    probabilistic and adaptive backtracking loops.
+
+    Returns a dict of everything both loops need. `state` is already replicated.
+    """
+    K = int(config.num_latent_slots)
+    assert K > 0, "Backtracking is for the recurrent latent model (K>0)."
+
+    logging.info("Creating datasets (backtracking loop)")
+    curriculum = data.CurriculumState(
+        stage=int(getattr(config, "curriculum_start_stage", 1)),
+        max_stage=int(getattr(config, "curriculum_max_stage", 6)))
+    # Per-level train sampler (replaces the uniform-over-unlocked sampler so we
+    # can draw a batch of a specific replay stage's level on demand).
+    sampler = StageBatchSampler(config, seed=int(config.seed))
+    eval_data_iter = data.create_iter(config, config.minibatch_size, train=False)
+
+    model_config = model.TransformerConfig(
+        dtype=config.dtype, vocab_size=config.vocab_size, seq_len=config.seq_len,
+        num_heads=config.num_heads, num_layers=config.num_layers,
+        emb_dim=config.emb_dim, qkv_dim=config.qkv_dim, mlp_dim=config.mlp_dim,
+        dropout_rate=config.dropout_rate,
+        attention_dropout_rate=config.attention_dropout_rate,
+        deterministic=False, num_latent_slots=K,
+        inject_latents=bool(int(getattr(config, "recurrent_latent", 1))),
+    )
+    print(str(model_config.__dict__), flush=True)
+
+    rng = jax.random.PRNGKey(config.seed)
+    rng, init_rng, dropout_rng = jax.random.split(rng, 3)
+    net = model.TransformerLMHeadModel(model_config)
+    dummy_latents = jax.numpy.zeros(
+        (config.minibatch_size, K, config.emb_dim), model_config.dtype)
+    dummy_positions = jax.numpy.zeros((config.minibatch_size, K), jax.numpy.int32)
+    dummy_active = jax.numpy.zeros((config.minibatch_size, K), bool)
+    _, initial_variables = jax.jit(net.init_with_output)(
+        {"params": init_rng, "dropout": dropout_rng},
+        jax.numpy.ones((config.minibatch_size, config.seq_len), jax.numpy.int32),
+        dummy_latents, dummy_positions, dummy_active)
+
+    state, lr_scheduler_fn = trainer.get_state(config, net, initial_variables)
+    start_step = 0
+    if config.resume_training:
+        state = checkpoints.restore_checkpoint(config.ckpt_loc, state)
+        start_step = int(state.step)
+        print("----------Restored model from", config.ckpt_loc,
+              f"at step {start_step}-----------")
+
+    writer = metric_writers.create_default_writer(
+        workdir, asynchronous=False, just_logging=(jax.process_index() > 0))
+    tf_summary_writer = tf.summary.create_file_writer(workdir)
+
+    state = jax_utils.replicate(state)
+    dropout_rngs = jax.random.split(rng, jax.local_device_count())
+
+    def make_p_train_step(num_passes, backtrack):
+        return jax.pmap(
+            functools.partial(
+                trainer.train_step, config=model_config, hyperparams=config,
+                learning_rate_fn=lr_scheduler_fn, num_passes=num_passes,
+                backtrack=backtrack),
+            axis_name="batch", donate_argnums=(0,))
+
+    # Precompile, for every recurrence depth 1..K, both a FRONTIER step (normal
+    # loss: LM CE + all active-slot BCE) and a strict BACKTRACK step (only the
+    # last active slot's readout, no CE). Frontier is used when t == curriculum
+    # stage i; the strict backtrack step is used by the probabilistic loop for
+    # replayed earlier stages t < i.
+    # Stage -> recurrence depth. With passes_per_stage>1 a stage advances the
+    # chain by more than one slot (pps=2, K=12: stage 6 -> all 12 slots), so the
+    # replay depth for stage t is pps*t, not t. Keyed by stage either way.
+    pps = int(getattr(config, "passes_per_stage", 1))
+    _depth = lambda t: min(pps * t, K)
+    p_frontier = {t: make_p_train_step(_depth(t), False) for t in range(1, K + 1)}
+    p_backtrack = {t: make_p_train_step(_depth(t), True) for t in range(1, K + 1)}
+    p_eval_step = jax.pmap(functools.partial(
+        evaluater.eval_step, config=model_config.replace(deterministic=True)),
+        axis_name="batch")
+
+    hooks, report_progress, _ = trainer.get_metrics_report_progress(
+        config, workdir, writer)
+
+    return {
+        "K": K,
+        "start_step": start_step,
+        "curriculum": curriculum,
+        "sampler": sampler,
+        "eval_data_iter": eval_data_iter,
+        "model_config": model_config,
+        "state": state,
+        "dropout_rngs": dropout_rngs,
+        "p_frontier": p_frontier,
+        "p_backtrack": p_backtrack,
+        "p_eval_step": p_eval_step,
+        "hooks": hooks,
+        "report_progress": report_progress,
+        "writer": writer,
+        "tf_summary_writer": tf_summary_writer,
+        "promote_threshold": float(getattr(config, "promote_acc_threshold", 0.85)),
+        "promote_patience": int(getattr(config, "promote_patience_steps", 8000)),
+        "min_stage_steps": int(getattr(config, "min_stage_steps", 2000)),
+        "ckpt_keep": int(getattr(config, "ckpt_keep", 100)),
+        "stage_ckpt_dir": os.path.join(workdir, "stage_ckpts"),
+    }
+
+
+def train_and_evaluate_backtrack(config, workdir):
+    """Backtracking curriculum training loop (stage replay with matched depth)."""
+    workdir = os.path.abspath(workdir)
+    backtrack_prob = float(getattr(config, "backtrack_prob", 0.5))
+    frontier_mix = bool(int(getattr(config, "backtrack_frontier_mix", 1)))
+
+    ctx = _prepare_backtrack(config, workdir)
+    K = ctx["K"]
+    curriculum = ctx["curriculum"]
+    sampler = ctx["sampler"]
+    eval_data_iter = ctx["eval_data_iter"]
+    state = ctx["state"]
+    dropout_rngs = ctx["dropout_rngs"]
+    p_frontier = ctx["p_frontier"]
+    p_backtrack = ctx["p_backtrack"]
+    p_eval_step = ctx["p_eval_step"]
+    hooks = ctx["hooks"]
+    writer = ctx["writer"]
+    tf_summary_writer = ctx["tf_summary_writer"]
+    promote_threshold = ctx["promote_threshold"]
+    promote_patience = ctx["promote_patience"]
+    min_stage_steps = ctx["min_stage_steps"]
+    ckpt_keep = ctx["ckpt_keep"]
+    stage_ckpt_dir = ctx["stage_ckpt_dir"]
+    stage_started_at = 0
+    sched_rng = np.random.RandomState(int(config.seed) + 1)
+
+    # Bookkeeping: how often each replay depth is actually trained.
+    stage_step_counts = {t: 0 for t in range(1, K + 1)}
+
+    def sample_target_stage(i):
+        """Frontier stage i with prob (1-p); else a backtrack stage 1..i-1."""
+        if i <= 1 or sched_rng.rand() >= backtrack_prob:
+            return i
+        return int(sched_rng.randint(1, i))   # uniform in {1..i-1}
+
+    with metric_writers.ensure_flushes(writer):
+        for step in range(0, config.max_steps):
+            if step % 10000 == 0:
+                print("Step:", step, flush=True)
+
+            i = curriculum.stage
+            t = sample_target_stage(i)
+            # Frontier step (t == i): normal loss. Backtrack step (t < i): strict
+            # replay of the depth-t readout only. Both draw from the whole
+            # corpus: depth, not difficulty, is the replayed axis.
+            step_fn = p_frontier[t] if t == i else p_backtrack[t]
+            batch = sampler.sample_all(config.minibatch_size)
+            state, metrics = _run_train_step(step_fn, state, batch, dropout_rngs)
+            stage_step_counts[t] += 1
+
+            for h in hooks:
+                h(step)
+
+            if math.isnan(metrics["loss"][0]):
+                print("Loss became nan; stopping.", flush=True)
+                break
+
+            if step % config.eval_every_steps == 0:
+                eval_metrics = evaluater.get_eval_metrics(
+                    state, eval_data_iter, p_eval_step, config)
+                per_level = eval_metrics.pop("per_level_acc")
+                per_depth = eval_metrics.pop("per_slot_acc_changed", {})
+                per_depth_all = eval_metrics.pop("per_slot_acc", {})
+                if not any(v >= 0 for v in per_depth.values()):
+                    per_depth = per_depth_all
+
+                def _m(key):
+                    v = eval_metrics.get(key, [])
+                    return round(float(np.mean(v)), 4) if len(v) else -1.0
+
+                print(step, "stage", curriculum.stage,
+                      "target_t", t,
+                      "loss", round(float(metrics["loss"].mean()), 4),
+                      "ce", round(float(metrics["ce_loss"].mean()), 4),
+                      "aux_bce", round(float(metrics["aux_loss"].mean()), 4),
+                      "| val_acc", _m("acc"), "loc_acc", _m("loc_acc"),
+                      "val|loc", _m("val_given_loc_acc"),
+                      "| cand_bit_acc", _m("cand_bit_acc"),
+                      "cand_set_acc", _m("cand_set_acc"),
+                      "cand_set_chg", _m("cand_set_acc_changed"),
+                      "| replay_counts", dict(stage_step_counts),
+                      flush=True)
+
+                with tf_summary_writer.as_default():
+                    tf.summary.scalar("loss", metrics["loss"].mean(), step=step)
+                    tf.summary.scalar("ce_loss", metrics["ce_loss"].mean(), step=step)
+                    tf.summary.scalar("aux_bce_loss", metrics["aux_loss"].mean(), step=step)
+                    tf.summary.scalar("curriculum_stage", curriculum.stage, step=step)
+                    for key in eval_metrics.keys():
+                        tf.summary.scalar("eval_" + key,
+                                          np.array(eval_metrics[key]).mean(), step=step)
+                    for lvl, v in per_level.items():
+                        if v >= 0:
+                            tf.summary.scalar(f"eval_acc_level_{lvl}", v, step=step)
+
+                # ---- Curriculum promotion (same rule as the standard loop) ----
+                if curriculum.stage < curriculum.max_stage:
+                    frontier_acc = per_depth.get(curriculum.stage, -1.0)
+                    steps_in_stage = step - stage_started_at
+                    hit_threshold = frontier_acc >= promote_threshold
+                    patience_over = steps_in_stage >= promote_patience
+                    if steps_in_stage >= min_stage_steps and (hit_threshold or patience_over):
+                        reason = "threshold" if hit_threshold else "patience"
+                        curriculum.stage += 1
+                        stage_started_at = step
+                        print(f"[curriculum] step {step}: promote to stage "
+                              f"{curriculum.stage} ({reason}; graduated depth "
+                              f"{curriculum.stage - 1} cand-set "
+                              f"acc={frontier_acc:.3f}); "
+                              f"backtrack pool now depths 1..{curriculum.stage-1}",
+                              flush=True)
+                        if config.save_checkpoint:
+                            unrep = jax_utils.unreplicate(state)
+                            checkpoints.save_checkpoint_multiprocess(
+                                workdir, unrep, step, keep=ckpt_keep, overwrite=True)
+                            checkpoints.save_checkpoint_multiprocess(
+                                stage_ckpt_dir, unrep, step, keep=100,
+                                overwrite=True, prefix=f"stage{curriculum.stage}_")
+
+            if config.save_checkpoint and step > 0 and step % config.save_every_steps == 0:
+                checkpoints.save_checkpoint_multiprocess(
+                    workdir, jax_utils.unreplicate(state), step,
+                    keep=ckpt_keep, overwrite=True)
+
+        if config.save_checkpoint:
+            checkpoints.save_checkpoint_multiprocess(
+                workdir, jax_utils.unreplicate(state), config.max_steps,
+                keep=ckpt_keep, overwrite=True)
+
+
+def train_and_evaluate_backtrack_adaptive(config, workdir):
+    """Deficit-driven ("adaptive") backtracking loop.
+
+    Difference from the probabilistic loop
+    --------------------------------------
+    The probabilistic loop replays a *uniformly random* earlier stage every step
+    with fixed probability p, regardless of whether that stage needs help. This
+    loop instead *watches* each earlier stage's held-out accuracy and only goes
+    back to repair a stage when it has actually regressed:
+
+      * Reference: when the curriculum promotes past stage t, we record that
+        stage's frontier-level (t+2) val_acc as its "graduation" accuracy.
+      * Trigger (relative drop): while at frontier stage i, at every eval we scan
+        earlier stages 1..i-1; a stage t is *in deficit* if its current level-(t+2)
+        val_acc has fallen below (graduation_acc[t] - margin).
+      * Selection (most-deficient first): if any stage is in deficit we switch the
+        training target to the single most-regressed stage and train there.
+      * Repair step: a FULL frontier-style step at the matched depth (level t+2,
+        num_passes=t, LM CE + all-active-slot BCE) -- NOT the strict readout-only
+        backtrack step. The metric we are trying to restore is placement accuracy
+        (driven by the LM CE), so the repair objective must include it.
+      * Exit (recover-or-cap): stay on the stage until its val_acc climbs back
+        above (graduation_acc[t] - margin), or until a max-repair-steps cap fires
+        (so a stuck stage cannot stall the frontier forever). On exit we re-scan
+        and either move to the next most-deficient stage or return to the frontier.
+
+    Frontier promotion is paused while repairing, and the repair time is credited
+    back to the frontier stage's patience clock on return.
+
+    Depth-matched eval note: the evaluator scores a level-L puzzle with exactly
+    k=L-2 active latent slots, so per_level_acc[t+2] is a faithful measurement of
+    the stage-t readout -- a clean trigger signal with no extra instrumentation.
+    """
+    workdir = os.path.abspath(workdir)
+    margin = float(getattr(config, "backtrack_margin", 0.03))
+    max_repair_steps = int(getattr(config, "backtrack_max_repair_steps", 6000))
+    min_frontier_steps = int(getattr(config, "backtrack_min_frontier_steps", 0))
+    frontier_target_acc = float(
+        getattr(config, "backtrack_frontier_target_acc", 0.0))
+    grad_acc_seed_raw = str(getattr(config, "backtrack_grad_acc_seed", "") or "")
+    max_repair_fraction = float(
+        getattr(config, "backtrack_max_repair_fraction", 0.0))
+    grad_decay = float(getattr(config, "backtrack_grad_decay", 0.0))
+    freeze_after_step = int(getattr(config, "backtrack_freeze_after_step", 0))
+    frontier_mix = bool(int(getattr(config, "backtrack_frontier_mix", 1)))
+
+    ctx = _prepare_backtrack(config, workdir)
+    K = ctx["K"]
+    start_step = ctx["start_step"]
+    curriculum = ctx["curriculum"]
+    sampler = ctx["sampler"]
+    eval_data_iter = ctx["eval_data_iter"]
+    state = ctx["state"]
+    dropout_rngs = ctx["dropout_rngs"]
+    p_frontier = ctx["p_frontier"]
+    p_eval_step = ctx["p_eval_step"]
+    hooks = ctx["hooks"]
+    writer = ctx["writer"]
+    tf_summary_writer = ctx["tf_summary_writer"]
+    promote_threshold = ctx["promote_threshold"]
+    promote_loc_threshold = float(getattr(config, "promote_loc_threshold", 0.70))
+    instance_mode = bool(getattr(config, "instance_dir", None))
+    promote_patience = ctx["promote_patience"]
+    min_stage_steps = ctx["min_stage_steps"]
+    ckpt_keep = ctx["ckpt_keep"]
+    stage_ckpt_dir = ctx["stage_ckpt_dir"]
+
+    # Controller state.
+    grad_acc = {}                 # stage t -> level-(t+2) val_acc at graduation
+    plateau_steps = int(getattr(config, "plateau_steps", 0))
+    plateau_delta = float(getattr(config, "plateau_delta", 0.005))
+    stage_best_acc = -1.0
+    stage_best_step = start_step
+    mode = "frontier"             # "frontier" | "repair"
+    repair_stage = None           # stage currently being repaired
+    repair_started_at = 0         # step the current repair stage began (cap)
+    repair_episode_start = 0      # step we first left the frontier (clock credit)
+    stage_started_at = start_step  # step the current frontier stage began
+    last_repair_return_step = start_step  # for min-frontier cooldown
+    step_counts = {t: 0 for t in range(1, K + 1)}   # steps trained at each depth
+    repair_steps_total = 0            # for duty-cycle cap
+    print(f"[repair] knobs: margin={margin} max_repair_steps={max_repair_steps} "
+          f"min_frontier={min_frontier_steps} frontier_target={frontier_target_acc} "
+          f"max_repair_frac={max_repair_fraction} grad_decay={grad_decay} "
+          f"freeze_after={freeze_after_step} frontier_mix={frontier_mix}",
+          flush=True)
+    # When seeding from a mid-curriculum checkpoint (start stage > 1), the
+    # earlier stages graduated in the *source* run so we have no reference for
+    # them. Prefer an explicit SUDOKU_GRAD_ACC_SEED (true graduation refs);
+    # otherwise seed from the first eval (old behavior — can over-trigger).
+    seed_start_stage = int(getattr(config, "curriculum_start_stage", 1))
+    grad_acc_seeded = seed_start_stage <= 1
+    if (not grad_acc_seeded) and grad_acc_seed_raw.strip():
+        try:
+            vals = [float(x) for x in grad_acc_seed_raw.split(",") if x.strip()]
+            for s, v in enumerate(vals, start=1):
+                if s < seed_start_stage:
+                    grad_acc[s] = v
+            if grad_acc:
+                grad_acc_seeded = True
+                print(f"[repair] seeded graduation refs from env: "
+                      f"{dict((k, round(v, 3)) for k, v in grad_acc.items())}",
+                      flush=True)
+        except ValueError:
+            print(f"[repair] WARNING: bad SUDOKU_GRAD_ACC_SEED="
+                  f"{grad_acc_seed_raw!r}; falling back to first-eval seeding",
+                  flush=True)
+
+    def compute_deficits(frontier_stage, per_depth):
+        """{depth t: graduation_acc[t] - current_acc} over earlier graduated
+        depths whose current candidate-set accuracy is measured.
+
+        Keyed on reasoning depth, not difficulty level: depth t's accuracy is
+        how well wave snapshot t is predicted, which is exactly what stage t
+        taught. A drop there means that propagation block has been forgotten."""
+        d = {}
+        for t in range(1, frontier_stage):
+            if t in grad_acc and per_depth.get(t, -1.0) >= 0:
+                d[t] = grad_acc[t] - per_depth.get(t, -1.0)
+        return d
+
+    def effective_margin(per_depth, frontier_stage):
+        """Widen the repair trigger while the frontier itself is still weak."""
+        m = margin
+        if frontier_target_acc > 0:
+            f_acc = per_depth.get(frontier_stage, -1.0)
+            if 0.0 <= f_acc < frontier_target_acc:
+                # Only repair clearer regressions until the frontier is good.
+                m = max(m, margin + (frontier_target_acc - f_acc))
+        return m
+
+    def most_deficient(frontier_stage, per_depth, use_margin=None):
+        """Most-regressed depth whose drop exceeds the margin, else None."""
+        d = compute_deficits(frontier_stage, per_depth)
+        if not d:
+            return None
+        m = margin if use_margin is None else use_margin
+        t = max(d, key=d.get)
+        return t if d[t] > m else None
+
+    with metric_writers.ensure_flushes(writer):
+        for step in range(start_step, config.max_steps):
+            if step % 10000 == 0:
+                print("Step:", step, flush=True)
+
+            i = curriculum.stage
+            # Target depth: the frontier when training normally, else the depth
+            # we are repairing. Difficulty is never gated, so a repair differs
+            # from a frontier step only in the recurrence depth it trains at:
+            # replaying depth t re-supervises wave snapshots 1..t on the same
+            # full-corpus batch distribution.
+            t = repair_stage if mode == "repair" else i
+            batch = sampler.sample_all(config.minibatch_size)
+            state, metrics = _run_train_step(
+                p_frontier[t], state, batch, dropout_rngs)
+            step_counts[t] += 1
+            if mode == "repair":
+                repair_steps_total += 1
+
+            for h in hooks:
+                h(step)
+
+            if math.isnan(metrics["loss"][0]):
+                print("Loss became nan; stopping.", flush=True)
+                break
+
+            if step % config.eval_every_steps == 0:
+                eval_metrics = evaluater.get_eval_metrics(
+                    state, eval_data_iter, p_eval_step, config)
+                per_level = eval_metrics.pop("per_level_acc")
+                # Depth accuracy drives the controller. Changed-cells-only, so
+                # a slot that merely copies its predecessor scores zero credit.
+                per_depth = eval_metrics.pop("per_slot_acc_changed", {})
+                per_depth_all = eval_metrics.pop("per_slot_acc", {})
+                per_stage_inset = eval_metrics.pop("per_stage_inset_acc", {})
+                if instance_mode:
+                    # Repair if a stage's in-set rate falls behind its
+                    # graduation value. The candidate head is off.
+                    per_depth = per_stage_inset
+                elif not any(v >= 0 for v in per_depth.values()):
+                    per_depth = per_depth_all
+
+                def _m(key):
+                    v = eval_metrics.get(key, [])
+                    return round(float(np.mean(v)), 4) if len(v) else -1.0
+
+                # Seed graduation refs for stages inherited from a checkpoint.
+                if not grad_acc_seeded:
+                    for s in range(1, seed_start_stage):
+                        acc_s = per_depth.get(s, -1.0)
+                        if acc_s >= 0:
+                            grad_acc[s] = float(acc_s)
+                    grad_acc_seeded = True
+                    print(f"[repair] step {step}: seeded graduation refs from "
+                          f"resume: "
+                          f"{dict((k, round(v, 3)) for k, v in grad_acc.items())}",
+                          flush=True)
+
+                # Soft graduation refs: forgive chronic mild regression.
+                if grad_decay > 0 and grad_acc:
+                    for s in list(grad_acc.keys()):
+                        cur_s = per_depth.get(s, -1.0)
+                        if cur_s >= 0 and cur_s < grad_acc[s]:
+                            old = grad_acc[s]
+                            grad_acc[s] = (
+                                (1.0 - grad_decay) * grad_acc[s]
+                                + grad_decay * float(cur_s))
+                            if step % (config.eval_every_steps * 5) == 0:
+                                print(f"[repair] soft-grad depth {s}: "
+                                      f"{old:.3f}->{grad_acc[s]:.3f} "
+                                      f"(cur={cur_s:.3f})", flush=True)
+
+                frontier_acc = per_depth.get(curriculum.stage, -1.0)
+                eff_margin = effective_margin(per_depth, i)
+                deficits = compute_deficits(i, per_depth)
+                elapsed = max(1, step - start_step + 1)
+                repair_frac = repair_steps_total / float(elapsed)
+                bt_frozen = (
+                    freeze_after_step > 0 and step >= freeze_after_step)
+                print(step, "stage", curriculum.stage,
+                      "mode", mode,
+                      "repair_stage", repair_stage,
+                      "target_t", t,
+                      "loss", round(float(metrics["loss"].mean()), 4),
+                      "ce", round(float(metrics["ce_loss"].mean()), 4),
+                      "aux_bce", round(float(metrics["aux_loss"].mean()), 4),
+                      "| val_acc", _m("acc"), "loc_acc", _m("loc_acc"),
+                      "val|loc", _m("val_given_loc_acc"),
+                      "| frontier_acc", round(float(frontier_acc), 4),
+                      "eff_margin", round(float(eff_margin), 4),
+                      "| cand_bit_acc", _m("cand_bit_acc"),
+                      "cand_set_acc", _m("cand_set_acc"),
+                      "cand_set_chg", _m("cand_set_acc_changed"),
+                      "| grad_acc", {k: round(v, 3) for k, v in grad_acc.items()},
+                      "deficits", {k: round(v, 3) for k, v in deficits.items()},
+                      "step_counts", dict(step_counts),
+                      "repair_frac", round(repair_frac, 3),
+                      "bt_frozen", bt_frozen,
+                      flush=True)
+
+                with tf_summary_writer.as_default():
+                    tf.summary.scalar("loss", metrics["loss"].mean(), step=step)
+                    tf.summary.scalar("ce_loss", metrics["ce_loss"].mean(), step=step)
+                    tf.summary.scalar("aux_bce_loss", metrics["aux_loss"].mean(), step=step)
+                    tf.summary.scalar("curriculum_stage", curriculum.stage, step=step)
+                    tf.summary.scalar("repair_mode", 1 if mode == "repair" else 0, step=step)
+                    tf.summary.scalar("repair_stage", repair_stage or 0, step=step)
+                    if frontier_acc >= 0:
+                        tf.summary.scalar("frontier_depth_acc", frontier_acc, step=step)
+                    tf.summary.scalar("eff_repair_margin", eff_margin, step=step)
+                    for key in eval_metrics.keys():
+                        tf.summary.scalar("eval_" + key,
+                                          np.array(eval_metrics[key]).mean(), step=step)
+                    for lvl, v in per_level.items():
+                        if v >= 0:
+                            tf.summary.scalar(f"eval_acc_level_{lvl}", v, step=step)
+                    for s, v in per_depth_all.items():
+                        if v >= 0:
+                            tf.summary.scalar(f"eval_cand_depth_{s}", v, step=step)
+
+                def _save_stage_ckpt(tag):
+                    if config.save_checkpoint:
+                        unrep = jax_utils.unreplicate(state)
+                        checkpoints.save_checkpoint_multiprocess(
+                            workdir, unrep, step, keep=ckpt_keep, overwrite=True)
+                        checkpoints.save_checkpoint_multiprocess(
+                            stage_ckpt_dir, unrep, step, keep=100,
+                            overwrite=True, prefix=f"{tag}_")
+
+                # ---------------- Deficit-driven controller ----------------
+                # Freeze / duty-cycle: force frontier-only when budget exhausted.
+                duty_ok = (
+                    max_repair_fraction <= 0
+                    or repair_frac < max_repair_fraction)
+                if bt_frozen and mode == "repair":
+                    print(f"[repair] step {step}: freeze_after="
+                          f"{freeze_after_step}; leaving repair", flush=True)
+                    mode = "frontier"
+                    repair_stage = None
+                    last_repair_return_step = step
+                if (not duty_ok) and mode == "repair":
+                    print(f"[repair] step {step}: duty-cycle cap "
+                          f"(repair_frac={repair_frac:.3f}>="
+                          f"{max_repair_fraction}); return to frontier",
+                          flush=True)
+                    mode = "frontier"
+                    repair_stage = None
+                    last_repair_return_step = step
+
+                if mode == "frontier":
+                    cooldown_ok = (
+                        min_frontier_steps <= 0
+                        or (step - last_repair_return_step) >= min_frontier_steps)
+                    can_repair = (not bt_frozen) and duty_ok and cooldown_ok
+                    worst = (most_deficient(i, per_depth, use_margin=eff_margin)
+                             if can_repair else None)
+                    if worst is not None:
+                        mode = "repair"
+                        repair_stage = worst
+                        repair_started_at = step
+                        repair_episode_start = step
+                        print(f"[repair] step {step}: enter repair of depth "
+                              f"{worst} (snapshot {worst} acc="
+                              f"{per_depth.get(worst, -1.0):.3f} < grad "
+                              f"{grad_acc.get(worst, -1.0):.3f} - "
+                              f"eff_margin {eff_margin:.3f}; "
+                              f"frontier_acc={frontier_acc:.3f})",
+                              flush=True)
+                    elif bt_frozen and most_deficient(
+                            i, per_depth, use_margin=eff_margin) is not None:
+                        if step % (config.eval_every_steps * 5) == 0:
+                            print(f"[repair] step {step}: deficit present but "
+                                  f"BT frozen after {freeze_after_step}",
+                                  flush=True)
+                    elif (not duty_ok) and most_deficient(
+                            i, per_depth, use_margin=eff_margin) is not None:
+                        if step % (config.eval_every_steps * 5) == 0:
+                            print(f"[repair] step {step}: deficit present but "
+                                  f"duty-cycle cap "
+                                  f"(frac={repair_frac:.3f})",
+                                  flush=True)
+                    elif (not cooldown_ok) and most_deficient(
+                            i, per_depth, use_margin=eff_margin) is not None:
+                        print(f"[repair] step {step}: deficit present but "
+                              f"frontier cooldown "
+                              f"({step - last_repair_return_step}/"
+                              f"{min_frontier_steps}); staying on stage {i}",
+                              flush=True)
+                    elif curriculum.stage < curriculum.max_stage:
+                        # Normal promotion (same rule as the standard loop).
+                        steps_in_stage = step - stage_started_at
+                        # A negative accuracy means "not measured this eval": it
+                        # must not reset the plateau tracker nor satisfy the
+                        # threshold/plateau rule. Only patience fires unmeasured.
+                        measured = frontier_acc >= 0
+                        if measured and frontier_acc > stage_best_acc + plateau_delta:
+                            stage_best_acc = frontier_acc
+                            stage_best_step = step
+                        loc_now = _m("loc_acc")
+                        loc_ready = (not instance_mode) or loc_now >= promote_loc_threshold
+                        hit_threshold = (measured and loc_ready
+                                         and frontier_acc >= promote_threshold)
+                        stalled = (measured and loc_ready and plateau_steps > 0
+                                   and (step - stage_best_step) >= plateau_steps)
+                        patience_over = steps_in_stage >= promote_patience
+                        if steps_in_stage >= min_stage_steps and (
+                                hit_threshold or stalled or patience_over):
+                            reason = ("threshold" if hit_threshold
+                                      else "plateau" if stalled else "patience")
+                            # Record this stage's graduation accuracy BEFORE moving on.
+                            grad_acc[curriculum.stage] = float(frontier_acc)
+                            curriculum.stage += 1
+                            stage_started_at = step
+                            stage_best_acc = -1.0
+                            stage_best_step = step
+                            print(f"[curriculum] step {step}: promote to stage "
+                                  f"{curriculum.stage} ({reason}; graduated "
+                                  f"depth {curriculum.stage - 1} cand-set "
+                                  f"acc={frontier_acc:.3f})", flush=True)
+                            _save_stage_ckpt(f"stage{curriculum.stage}")
+                    elif (frontier_target_acc > 0
+                          and 0.0 <= frontier_acc < frontier_target_acc
+                          and step % (config.eval_every_steps * 5) == 0):
+                        print(f"[frontier] step {step}: depth {i} "
+                              f"acc={frontier_acc:.3f} "
+                              f"< target {frontier_target_acc:.3f}; "
+                              f"keeping frontier priority",
+                              flush=True)
+                else:  # mode == "repair"
+                    r = repair_stage
+                    cur = per_depth.get(r, -1.0)
+                    ref = grad_acc.get(r, -1.0)
+                    recovered = cur >= (ref - margin)
+                    capped = (step - repair_started_at) >= max_repair_steps
+                    if recovered or capped:
+                        why = "recovered" if recovered else "cap"
+                        print(f"[repair] step {step}: depth {r} done ({why}; "
+                              f"snapshot acc={cur:.3f} vs grad {ref:.3f})",
+                              flush=True)
+                        _save_stage_ckpt(f"repair{r}")
+                        # Re-scan: chain to the next most-deficient stage, or
+                        # return to the frontier and credit repair time back to
+                        # the frontier stage's patience clock.
+                        nxt = None
+                        if (not bt_frozen) and duty_ok:
+                            nxt = most_deficient(
+                                i, per_depth, use_margin=eff_margin)
+                        if nxt is not None:
+                            repair_stage = nxt
+                            repair_started_at = step
+                            print(f"[repair] step {step}: chain to stage {nxt}",
+                                  flush=True)
+                        else:
+                            mode = "frontier"
+                            repair_stage = None
+                            last_repair_return_step = step
+                            stage_started_at += (step - repair_episode_start)
+
+            if config.save_checkpoint and step > 0 and step % config.save_every_steps == 0:
+                checkpoints.save_checkpoint_multiprocess(
+                    workdir, jax_utils.unreplicate(state), step,
+                    keep=ckpt_keep, overwrite=True)
+
+        if config.save_checkpoint:
+            checkpoints.save_checkpoint_multiprocess(
+                workdir, jax_utils.unreplicate(state), config.max_steps,
+                keep=ckpt_keep, overwrite=True)

code/train/trainer.py

@@ -0,0 +1,383 @@
+# coding=utf-8
+# Copyright 2024 The Google Research Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""Transformer LM trainer."""
+
+import functools
+
+from clu import periodic_actions
+from flax.training import common_utils
+from flax.training import train_state
+import jax
+from jax import numpy as jnp
+import numpy as np
+import optax
+import ml_collections
+
+from train import model
+
+
+def get_state(config, net, initial_variables):
+    """Get the train state given an experiment config, a model and initial variables.
+
+    Args:
+        config: A ConfigDict containing the configuration for the experiment.
+        net: The model to use for training.
+        initial_variables: The initial variables for the model.
+
+    Returns:
+        A tuple containing the train state and the learning rate schedule.
+    """
+    # Learning rate schedule
+    lr_scheduler_fn = functools.partial(
+        lr_scheduler,
+        learning_rate=config.learning_rate,
+        warmup_tokens=config.warmup_tokens,
+        final_tokens=config.max_steps,
+        config=config,
+    )
+    # Optimizer
+    optim_fn = optax.adamw(
+        lr_scheduler_fn, weight_decay=config.weight_decay, b1=0.9, b2=0.95
+    )
+    # Clip the gradients to prevent exploding gradients.
+    optimizer = optax.chain(optax.clip_by_global_norm(1), optim_fn)
+
+    # Initialize the train state
+    state = train_state.TrainState.create(
+        apply_fn=net.apply, params=initial_variables["params"],
+        tx=optimizer
+        )
+
+    return state, lr_scheduler_fn
+
+def lr_scheduler(
+    n_tokens: int, learning_rate: float, warmup_tokens: int, final_tokens: int,
+    config: ml_collections.ConfigDict,
+) -> float:
+    """Learning rate scheduler, adapted from Mikhail Grankin.
+
+    The learning rate schedule is cosine decay with a warmup period.
+
+    The learning rate starts at 0 and linearly increases to the given learning
+    rate over the warmup period. After the warmup period, the learning rate
+    decays according to a cosine schedule, with the given learning rate as the
+    maximum value.
+
+    Args:
+        n_tokens: The number of tokens processed so far.
+        learning_rate: The initial learning rate.
+        warmup_tokens: The number of tokens to warm up over.
+        final_tokens: The total number of tokens to process.
+        config: A ConfigDict containing the configuration for the learning rate
+            schedule.
+
+    Returns:
+        The learning rate at the given point in the schedule.
+    """
+    # Decay the learning rate based on our progress.
+    progress = (n_tokens - warmup_tokens) / max(
+        1, final_tokens - warmup_tokens,
+    )
+    lr_mult = jnp.where(
+        n_tokens < warmup_tokens,
+        # Linear warmup.
+        n_tokens / jnp.fmax(1, warmup_tokens),
+        # Cosine learning rate decay.
+        jnp.fmax(config.end_lr_factor, 0.5 * (1.0 + jnp.cos(np.pi * progress))),
+    )
+    return learning_rate * lr_mult
+
+
+def get_metrics_report_progress(config, workdir, writer):
+    """
+    Get the metrics for reporting progress during training.
+
+    Args:
+        config: The configuration for the experiment.
+        workdir: The directory for storing the logs.
+        writer: The writer object for recording the metrics.
+
+    Returns:
+        hooks: List of hooks for tracking progress.
+        report_progress: Object for reporting progress.
+        train_metrics: List of training metrics.
+    """
+    hooks = []
+
+    # Initialize the report progress object
+    report_progress = periodic_actions.ReportProgress(
+        num_train_steps=config.max_steps, writer=writer)
+
+    # Add metrics for profiling if the process index is 0
+    if jax.process_index() == 0:
+        hooks += [report_progress,
+                  periodic_actions.Profile(logdir=workdir, num_profile_steps=5)]
+    
+    # Initialize the list of training metrics
+    train_metrics = []
+    
+    return hooks, report_progress, train_metrics
+
+
+def get_input_start_index(batch, config):
+    inputs = jax.tree_util.tree_map(np.asarray, batch[0])
+    puzzles = jax.tree_util.tree_map(np.asarray, batch[1])
+    start_index = jax.tree_util.tree_map(np.asarray, batch[2])
+    levels = jax.tree_util.tree_map(np.asarray, batch[3])
+    cand_targets = jax.tree_util.tree_map(np.asarray, batch[4])
+    return inputs, puzzles, start_index, levels, cand_targets
+
+
+def train_one_step(p_train_step, config, state, step, dropout_rngs, train_data_iter):
+    """
+    Single step of the training loop.
+
+    Args:
+        p_train_step: The training step function.
+        config: The experiment configuration.
+        state: The train state.
+        step: The step number.
+        dropout_rngs: The dropout random number generator.
+        train_data_iter: The iterator for the train data.
+
+    Returns:
+        The updated train state and train metrics.
+    """
+    with jax.profiler.StepTraceAnnotation("train", step_num=step):
+        # Get the next batch from the iterator
+        batch = next(train_data_iter)
+        # Extract the inputs, start index and difficulty level from the batch
+        inputs, _, start_index, levels, cand_targets = get_input_start_index(batch, config)
+        # Shard across the devices
+        inputs = common_utils.shard(jax.tree_util.tree_map(np.asarray, inputs))
+        start_index = common_utils.shard(jax.tree_util.tree_map(np.asarray, start_index))
+        levels = common_utils.shard(jax.tree_util.tree_map(np.asarray, levels))
+        # Explicit reshape (not common_utils.shard) so a zero-width slot dim
+        # (K=0 no-latent baseline: cand_targets is (bs, 0, 81)) shards without
+        # the ambiguous -1 inference. Identical to shard() when K>0.
+        cand_targets = np.asarray(cand_targets)
+        _nd = jax.local_device_count()
+        cand_targets = cand_targets.reshape(
+            (_nd, cand_targets.shape[0] // _nd) + cand_targets.shape[1:])
+
+        # Run the training step
+        state, metrics, _ = p_train_step(
+            state, inputs, start_index, levels, cand_targets, dropout_rng=dropout_rngs
+        )
+
+    return state, metrics
+
+
+def build_latent_state(inputs, start_index, levels, config, num_passes,
+                       apply_fn, rngs=None):
+    """Build the continuous latent thoughts (Coconut/ATC-style recurrence).
+
+    Pass 0 seeds z_1 from the last-layer hidden at the last clue token; pass j
+    reads the hidden at latent slot j to produce z_{j+1}. Slot j is active for
+    an example only when j < k, with k = clip(level - 2, 1, K) (difficulty-
+    matched latent budget). Gradients flow through all passes (full BPTT).
+
+    Returns (latent_vals, slot_pos, active_full).
+    """
+    num_slots = int(config.num_latent_slots)
+    bs = inputs.shape[0]
+    bidx = jnp.arange(bs)
+
+    si3 = 3 * start_index.reshape(-1)                      # (bs,)
+    slot_pos = si3[:, None] + jnp.arange(num_slots)[None, :]   # (bs, K)
+    if getattr(config, "cand_slot_mode", "level") == "depth":
+        # Every slot the recurrence fills is active, uniformly over the batch.
+        k = jnp.full((bs,), max(min(num_passes, num_slots), 1), dtype=jnp.int32)
+    else:
+        k = jnp.clip(levels.reshape(-1) - 2, 1, num_slots)  # (bs,)
+    active_full = jnp.arange(num_slots)[None, :] < k[:, None]  # (bs, K)
+
+    latent_vals = jnp.zeros((bs, num_slots, config.emb_dim), dtype=config.dtype)
+
+    for j in range(num_passes):
+        act_j = active_full & (jnp.arange(num_slots)[None, :] < j)
+        _, hidden, _ = apply_fn(inputs, latent_vals, slot_pos, act_j, rngs)
+        src = si3 - 1 + j   # j=0: last clue token; j>0: latent slot j-1
+        z = hidden[bidx, src]
+        latent_vals = latent_vals.at[:, j].set(z.astype(config.dtype))
+
+    return latent_vals, slot_pos, active_full
+
+
+def train_step(state, batch, start_index, levels, cand_targets, config,
+               hyperparams, learning_rate_fn, num_passes, dropout_rng=None,
+               backtrack=False):
+    """One step of the training loop.
+
+    Args:
+        state: Train state.
+        batch: Input batch (bs, 3*81 + K) with latent placeholder slots.
+        start_index: Number of clue cells per example.
+        levels: Puzzle difficulty level (3..8) per example.
+        config: Model config.
+        hyperparams: Hyperparameter dictionary.
+        learning_rate_fn: Learning rate function.
+        num_passes: Number of latent recurrence passes (= curriculum stage).
+        dropout_rng: RNG used for dropout.
+
+    Returns:
+        A new train state, train metrics, and computed model predictions.
+    """
+    num_slots = int(config.num_latent_slots)
+    # Extract inputs and labels from the batch
+    inputs = batch[:, :-1]
+    label = batch[:, 1:]
+
+    # Update dropout_rng
+    dropout_rng = jax.random.fold_in(dropout_rng, state.step)
+    dropout_rng_dict = {"dropout": dropout_rng}
+
+    def loss_fn(params):
+        """Compute the loss function."""
+        net = model.TransformerLMHeadModel(config)
+
+        def apply_fn(x, lv, lp, la, rngs):
+            return net.apply({"params": params}, x, latent_values=lv,
+                             latent_positions=lp, latent_active=la,
+                             rngs=rngs)
+
+        latent_vals, slot_pos, active_full = build_latent_state(
+            inputs, start_index, levels, config, num_passes, apply_fn,
+            rngs=dropout_rng_dict)
+
+        pred_logits, _, cand_logits = apply_fn(
+            inputs, latent_vals, slot_pos, active_full, dropout_rng_dict)
+
+        label_one_hot = jax.nn.one_hot(label, num_classes=config.vocab_size)
+
+        # The variables label_one_hot and pred_logits both are 3-dimensional tensors with
+        # first axis corresponding to batch size, second correspondingn to sequence length 
+        # and third corresponding to the row/column/value at a particular cell
+        assert label_one_hot.shape == pred_logits.shape, ("one hot label shape",
+                                                        label_one_hot.shape,
+                                                        label.shape,
+                                                        pred_logits.shape)
+        
+        # Calculate the cross-entropy loss along the last axis
+        pred_logits_sol = pred_logits[:, :, :]
+        label_one_hot_sol = label_one_hot[:, :, :]
+
+        ce_loss = optax.softmax_cross_entropy(
+            logits=pred_logits_sol[:, :, :], labels=label_one_hot_sol[:, :, :]
+        )
+        # assert ce_loss.ndim == 2, ("ce_loss", ce_loss.shape)
+
+        # Apply masking to the loss: supervise only the solution region,
+        # which now starts K latent slots after the clue block.
+        mask = np.repeat(
+            np.arange(len(ce_loss[0])).reshape(1, -1), len(ce_loss), axis=0
+        )
+        mask = (mask >= 3 * start_index + num_slots)      
+
+        # Per-token mean (not per-example sum): keeps the LM CE on the same
+        # O(1) scale as the per-digit BCE below, so aux_cand_weight~1 actually
+        # balances the two instead of the candidate signal being swamped.
+        ce_denom = jnp.maximum(mask.sum(), 1.0)
+        avg_ce_loss = (ce_loss * mask).sum() / ce_denom
+
+        # ---- Auxiliary multi-candidate BCE loss on the latent slots ----
+        # cand_targets: (bs, K, 81) int bitmask (bit d-1 set <=> digit d is a
+        # candidate at that slot's stage; slot j already aligned to the puzzle's
+        # k=level-2 budget in the data pipeline). Expand to (bs,K,81,9) multi-hot.
+        aux_weight = float(getattr(hyperparams, "aux_cand_weight", 1.0))
+        if aux_weight == 0.0:
+            # Instance arm: the superposition lives in the varying targets, not
+            # in a set head. Drop the whole BCE graph so nothing but the LM CE
+            # shapes the latents.
+            zero = jnp.zeros((), dtype=pred_logits.dtype)
+            return avg_ce_loss, (pred_logits, avg_ce_loss, zero)
+
+        bits = jnp.arange(9)
+        cand_multi_hot = ((cand_targets[..., None].astype(jnp.int32)
+                           >> bits) & 1).astype(cand_logits.dtype)  # (bs,K,81,9)
+        # Positive-weighted BCE. Candidate masks are sparse (~1-3 of 9 digits
+        # "on"), so plain BCE collapses to predicting all-zeros. Up-weighting
+        # the positive (candidate-present) term by pos_weight counteracts the
+        # imbalance and forces the head to actually predict the candidate set.
+        pos_weight = float(getattr(hyperparams, "aux_pos_weight", 5.0))
+        log_p = jax.nn.log_sigmoid(cand_logits)                     # log sigmoid(x)
+        log_1mp = jax.nn.log_sigmoid(-cand_logits)                  # log(1 - sigmoid(x))
+        bce = -(pos_weight * cand_multi_hot * log_p
+                + (1.0 - cand_multi_hot) * log_1mp)                 # (bs,K,81,9)
+        # Only the empty cells are supervised: clue cells were sentinel-zeroed
+        # in the data pipeline, so any cell whose target row is all-zero is a
+        # clue and must not contribute. This makes the effective target
+        # (#empty cells) x 9 per puzzle rather than 81 x 9.
+        cell_mask = (cand_targets > 0).astype(cand_logits.dtype)    # (bs, K, 81)
+        # Delta weighting: consecutive stages are near-duplicates (at K=12 only
+        # ~14% of cells change per stage, so ~97% of the target bits are copies
+        # of the previous slot). Weighting the unchanged cells below 1 stops the
+        # objective from being satisfied by echoing slot j-1. Slot 0 has no
+        # predecessor, so it is fully weighted.
+        delta_bg = float(getattr(hyperparams, "aux_delta_bg", 1.0))
+        if delta_bg != 1.0:
+            changed = jnp.concatenate(
+                [jnp.ones_like(cand_targets[:, :1], dtype=bool),
+                 cand_targets[:, 1:] != cand_targets[:, :-1]], axis=1)
+            cell_w = cell_mask * (delta_bg + (1.0 - delta_bg)
+                                  * changed.astype(cand_logits.dtype))
+        else:
+            cell_w = cell_mask
+        cell_denom = jnp.maximum(cell_w.sum(axis=2), 1e-6)          # (bs, K)
+        bce_per_cell = bce.mean(axis=3)                             # (bs, K, 81), mean over 9 digits
+        bce_per_slot = ((bce_per_cell * cell_w).sum(axis=2)
+                        / cell_denom)                               # (bs, K), mean over empty cells
+        if backtrack:
+            # Strict stage-replay: supervise ONLY the stage-t readout = the last
+            # active latent slot per example (index k-1), whose candidate target
+            # is that stage's grid. Earlier slots and the LM CE are dropped, so
+            # this step purely re-derives "apply f exactly t times -> stage-t
+            # readout." Injection above still spans all active slots, so the
+            # recurrence reaching slot k-1 is intact.
+            k_per = active_full.sum(axis=1)                         # (bs,)
+            last_idx = (k_per - 1)[:, None]                         # (bs, 1)
+            slot_sel = ((jnp.arange(num_slots)[None, :] == last_idx)
+                        & active_full)                             # (bs, K)
+            slot_active = slot_sel.astype(bce_per_slot.dtype)
+        else:
+            slot_active = active_full.astype(bce_per_slot.dtype)    # (bs, K)
+        aux_denom = jnp.maximum(slot_active.sum(), 1.0)
+        avg_aux_loss = (bce_per_slot * slot_active).sum() / aux_denom
+
+        # Backtrack (replay) steps train only from the readout (no LM CE).
+        ce_term = 0.0 if backtrack else avg_ce_loss
+        total_loss = ce_term + aux_weight * avg_aux_loss
+
+        return total_loss, (pred_logits, avg_ce_loss, avg_aux_loss)
+
+    # Compute the learning rate and perform gradient descent
+    step = state.step
+    lr = learning_rate_fn(step)
+    (loss, aux), grads = jax.value_and_grad(loss_fn,
+                                            has_aux=True)(state.params)
+    pred_logits, ce_loss, aux_loss = aux
+    grads = jax.lax.pmean(grads, "batch")
+    new_state = state.apply_gradients(grads=grads)
+    
+    # Update training metrics
+    metrics = {
+        "step": step, "loss": loss, "learning_rate": lr,
+        "ce_loss": ce_loss,
+        "aux_loss": aux_loss,
+        "pred_logits": pred_logits, "weights": inputs.shape[0]
+    }
+    
+    return new_state, metrics, pred_logits

code/verify_superposition_instances.py

@@ -0,0 +1,98 @@
+"""Statistics for superposition-instance generation across all stages.
+
+Reports, per stage: how many instances are produced per puzzle, how much of each
+candidate set the instances cover, how many proposals get ruled out by the
+dependencies, and how the dependency load changes as the grid resolves.
+"""
+import argparse
+import collections
+import multiprocessing as mp
+
+import numpy as np
+
+import superposition_instances as SI
+
+_MASKS = None
+_ARGS = None
+
+
+def _init(path, args):
+    global _MASKS, _ARGS
+    _MASKS = np.load(path, mmap_mode="r")
+    _ARGS = args
+
+
+def _one(idx):
+    out = []
+    n_stages = _MASKS.shape[1]
+    for s in range(n_stages):
+        mask = np.array(_MASKS[idx, s]).astype(np.uint16)
+        r = SI.instances_for_stage(
+            mask, max_confine=_ARGS.max_confine,
+            max_instances=_ARGS.max_instances,
+            max_attempts=_ARGS.max_attempts, seed=idx * 100 + s,
+            max_repair=_ARGS.max_repair)
+        all_deps = SI.dependencies_from_mask(mask, max_confine=9)
+        r["n_deps_all"] = len(all_deps)
+        r.pop("instances")
+        r["stage"] = s
+        out.append(r)
+    return out
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--masks", default="datasets_multicandidate_s12/test_cand_masks.npy")
+    ap.add_argument("--limit", type=int, default=200)
+    ap.add_argument("--max-confine", type=int, default=3)
+    ap.add_argument("--max-instances", type=int, default=32)
+    ap.add_argument("--max-attempts", type=int, default=200)
+    ap.add_argument("--max-repair", type=int, default=120)
+    ap.add_argument("--workers", type=int, default=32)
+    args = ap.parse_args()
+
+    masks = np.load(args.masks, mmap_mode="r")
+    n = min(args.limit, len(masks))
+    print(f"masks {masks.shape} from {args.masks}")
+    print(f"puzzles {n}, confinement threshold |S| <= {args.max_confine}, "
+          f"instance cap {args.max_instances}, attempt cap {args.max_attempts}")
+
+    with mp.Pool(args.workers, initializer=_init,
+                 initargs=(args.masks, args)) as pool:
+        results = pool.map(_one, range(n), chunksize=4)
+
+    by_stage = collections.defaultdict(list)
+    for rows in results:
+        for r in rows:
+            by_stage[r["stage"]].append(r)
+
+    print()
+    print("=" * 112)
+    print("PER-STAGE INSTANCE GENERATION")
+    print("=" * 112)
+    print(f"{'stage':>5} {'empty':>6} {'width':>6} {'deps all':>9} "
+          f"{'deps enf':>9} {'kept':>6} {'ruled out':>10} {'reject %':>9} "
+          f"{'coverage':>9} {'capped %':>9} {'dup/inst':>9}")
+    print("-" * 112)
+    for s in sorted(by_stage):
+        rs = by_stage[s]
+        g = lambda k: np.mean([r[k] for r in rs])
+        rej = np.sum([r["n_rejected"] for r in rs])
+        att = np.sum([r["n_attempts"] for r in rs])
+        print(f"{s:>5} {g('n_empty'):>6.1f} {g('mean_width'):>6.2f} "
+              f"{g('n_deps_all'):>9.1f} {g('n_deps'):>9.1f} "
+              f"{g('n_instances'):>6.1f} {g('n_rejected'):>10.1f} "
+              f"{100*rej/max(att,1):>8.1f}% {100*g('coverage'):>8.1f}% "
+              f"{100*np.mean([r['hit_cap'] for r in rs]):>8.1f}% "
+              f"{g('mean_dups'):>9.2f}")
+
+    tot_inst = np.sum([r["n_instances"] for rs in by_stage.values() for r in rs])
+    tot_rej = np.sum([r["n_rejected"] for rs in by_stage.values() for r in rs])
+    print()
+    print(f"per puzzle across all {len(by_stage)} stages: "
+          f"{tot_inst/n:.1f} instances kept, {tot_rej/n:.1f} ruled out")
+    print(f"dataset multiplier vs one row per puzzle: {tot_inst/n:.1f}x")
+
+
+if __name__ == "__main__":
+    main()