sync training code: stage-1 instance-epoch sampler, multi-stage run, superposition metrics

code/wavecurriculum_run/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/wavecurriculum_run/hf_push_login.py

@@ -0,0 +1,191 @@
+#!/usr/bin/env python3
+"""Upload scratch/gandalf staging files to Hugging Face.
+
+Runs on the login node (python3 + huggingface_hub). The feanor sidecar only
+needs to land files in --stage-dir; this process does the HTTP/LFS upload.
+"""
+import argparse
+import json
+import os
+import sys
+import time
+import traceback
+
+
+def _load_json(path):
+    try:
+        with open(path) as f:
+            return json.load(f)
+    except Exception:
+        return {}
+
+
+def _write_json(path, payload):
+    tmp = path + ".tmp"
+    with open(tmp, "w") as f:
+        json.dump(payload, f, indent=2, sort_keys=True)
+        f.write("\n")
+    os.replace(tmp, path)
+
+
+_KEEP_HB = (
+    "event", "step", "stage", "loss",
+    "val_acc", "val_mass", "val_excess", "val_spread",
+    "loc_wave", "loc_coverage", "loc_lcs", "loc_dup",
+)
+
+
+def _slim_heartbeat(src, dest):
+    raw = _load_json(src)
+    slim = {k: raw[k] for k in _KEEP_HB if k in raw}
+    _write_json(dest, slim)
+    return dest
+
+
+def _slim_log(src, dest):
+    """Keep only eval / promote lines: step, stage, loss, accuracies."""
+    keep = []
+    try:
+        with open(src, errors="replace") as f:
+            for line in f:
+                s = line.strip()
+                if not s:
+                    continue
+                if s.startswith("[curriculum]"):
+                    keep.append(s + "\n")
+                    continue
+                # "2000 stage 1 loss 1.47 loc_acc 0.03 val_acc 0.11 inset 0.66"
+                if " loss " in s and "stage" in s and not s.startswith("I"):
+                    keep.append(s + "\n")
+    except OSError:
+        return None
+    with open(dest, "w") as f:
+        f.writelines(keep)
+    return dest
+
+
+def _upload(api, token, local, repo, dest):
+    api.upload_file(
+        path_or_fileobj=local,
+        path_in_repo=dest,
+        repo_id=repo,
+        repo_type="model",
+        token=token,
+        commit_message=f"sync {dest}",
+    )
+    print(f"[hf_push] uploaded {dest} ({os.path.getsize(local)} bytes)", flush=True)
+
+
+def sync_once(args, state, api, token):
+    prefix = args.prefix.strip("/")
+    repo = args.repo
+
+    for local, dest, key in (
+        (args.heartbeat, f"{prefix}/heartbeat.json", "heartbeat"),
+        (args.log, f"{prefix}/train.log", "log"),
+    ):
+        if not local or not os.path.isfile(local):
+            continue
+        mtime = os.path.getmtime(local)
+        size = os.path.getsize(local)
+        last = state.get(key, {})
+        min_age = 0 if key == "heartbeat" else 300
+        due = (time.time() - float(last.get("uploaded_at", 0))) >= min_age
+        if due and (mtime > float(last.get("mtime", 0)) or size != last.get("size", -1)):
+            upload_path = local
+            if key == "heartbeat":
+                upload_path = _slim_heartbeat(
+                    local, os.path.join(os.path.dirname(local) or ".",
+                                        "_hf_heartbeat_slim.json"))
+            elif key == "log":
+                upload_path = _slim_log(
+                    local, os.path.join(os.path.dirname(local) or ".",
+                                        "_hf_train_slim.log"))
+            if upload_path:
+                _upload(api, token, upload_path, repo, dest)
+            state[key] = {"mtime": mtime, "size": size, "uploaded_at": time.time()}
+
+    if args.stage_dir and os.path.isdir(args.stage_dir):
+        for name in sorted(os.listdir(args.stage_dir)):
+            if name in ("hf_push_state.json", "heartbeat.json"):
+                continue
+            if not name.endswith(".tar") and name not in ("latest.json",):
+                continue
+            local = os.path.join(args.stage_dir, name)
+            if not os.path.isfile(local):
+                continue
+            # step_/stage_ names are unique per checkpoint, so once uploaded
+            # they never need revisiting. latest.tar/latest.json keep the same
+            # name and are rewritten every save, so skipping them by name would
+            # freeze the rolling checkpoint at whichever step happened to land
+            # first; track those by mtime instead.
+            rolling = name in ("latest.tar", "latest.json")
+            mtime = os.path.getmtime(local)
+            if rolling:
+                if mtime <= float(state.get("rolling", {}).get(name, 0)):
+                    continue
+            elif name in state.get("staged", []):
+                continue
+            dest = f"{prefix}/{name}"
+            if name.startswith("stage"):
+                dest = f"{prefix}/stages/{name}"
+            elif name.startswith("step_"):
+                dest = f"{prefix}/steps/{name}"
+            _upload(api, token, local, repo, dest)
+            if rolling:
+                state.setdefault("rolling", {})[name] = mtime
+                continue      # keep the file; the sidecar overwrites it in place
+            state.setdefault("staged", []).append(name)
+            # Drop the local copy after a successful upload to save quota.
+            try:
+                os.remove(local)
+            except OSError:
+                pass
+    return state
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--repo", default="Avra98/Sudoku_superposition")
+    ap.add_argument("--prefix", default="runs/w12_inst_latent")
+    ap.add_argument("--heartbeat", default="")
+    ap.add_argument("--log", default="")
+    ap.add_argument("--stage-dir", default="")
+    ap.add_argument("--token-file", default="/scratch/users/gatmiry/.hf_token")
+    ap.add_argument("--state", default="")
+    ap.add_argument("--interval", type=int, default=60)
+    ap.add_argument("--once", action="store_true")
+    args = ap.parse_args()
+
+    token = ""
+    if args.token_file and os.path.isfile(args.token_file):
+        with open(args.token_file) as f:
+            token = f.read().strip()
+    token = token or os.environ.get("HF_TOKEN", "")
+    if not token:
+        print("[hf_push] no token", file=sys.stderr)
+        return 1
+
+    from huggingface_hub import HfApi
+    api = HfApi(token=token)
+    state_path = args.state or os.path.join(
+        args.stage_dir or "/tmp/sudoku_hf_uploads", "hf_push_state.json")
+    os.makedirs(os.path.dirname(state_path) or ".", exist_ok=True)
+    state = _load_json(state_path)
+    print(f"[hf_push] watching heartbeat={args.heartbeat} log={args.log} "
+          f"stage={args.stage_dir}", flush=True)
+    while True:
+        try:
+            state = sync_once(args, state, api, token)
+            _write_json(state_path, state)
+        except Exception:
+            traceback.print_exc()
+            print("[hf_push] cycle failed; will retry", flush=True)
+        if args.once:
+            break
+        time.sleep(max(5, args.interval))
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

code/wavecurriculum_run/hf_sync.py

@@ -0,0 +1,253 @@
+#!/usr/bin/env python3
+"""Sidecar: copy logs to scratch and push heartbeat/checkpoints to Hugging Face.
+
+Does not train. Safe to run in a loop next to the job. Large Orbax dirs are
+tarred, then uploaded as a single file so `latest.tar` can be overwritten.
+"""
+import argparse
+import json
+import os
+import socket
+import subprocess
+import sys
+import time
+import traceback
+
+
+def _load_json(path):
+    try:
+        with open(path) as f:
+            return json.load(f)
+    except Exception:
+        return {}
+
+
+def _write_json(path, payload):
+    tmp = path + ".tmp"
+    with open(tmp, "w") as f:
+        json.dump(payload, f, indent=2, sort_keys=True)
+        f.write("\n")
+    os.replace(tmp, path)
+
+
+def _latest_ckpt_dir(workdir):
+    best = None
+    best_step = -1
+    try:
+        names = os.listdir(workdir)
+    except FileNotFoundError:
+        return None, -1
+    for name in names:
+        if not name.startswith("checkpoint_"):
+            continue
+        if name.endswith(".orbax-checkpoint-tmp") or ".orbax-checkpoint-tmp" in name:
+            continue
+        path = os.path.join(workdir, name)
+        if not os.path.isdir(path):
+            continue
+        try:
+            step = int(name.split("_", 1)[1])
+        except ValueError:
+            continue
+        if step > best_step:
+            best_step = step
+            best = path
+    return best, best_step
+
+
+def _stage_ckpt_dirs(stage_dir):
+    out = []
+    if not os.path.isdir(stage_dir):
+        return out
+    for name in sorted(os.listdir(stage_dir)):
+        path = os.path.join(stage_dir, name)
+        if os.path.isdir(path) and not name.endswith(".tmp"):
+            out.append(path)
+    return out
+
+
+def _tar(src_dir, dest_tar):
+    parent = os.path.dirname(os.path.abspath(src_dir))
+    base = os.path.basename(src_dir.rstrip("/"))
+    tmp = dest_tar + ".tmp"
+    subprocess.check_call(
+        ["tar", "-C", parent, "-cf", tmp, base],
+        stdout=subprocess.DEVNULL,
+        stderr=subprocess.DEVNULL,
+    )
+    os.replace(tmp, dest_tar)
+    return dest_tar
+
+
+def _hf_api(token):
+    try:
+        from huggingface_hub import HfApi
+    except Exception as exc:
+        print(f"[hf_sync] huggingface_hub unavailable ({exc}); "
+              f"will stage files for the login-node pusher", flush=True)
+        return None
+    return HfApi(token=token)
+
+
+def _rsync_to_login(local, remote_dir):
+    """Copy a file to gandalf so the login-node pusher can upload it."""
+    if not remote_dir or not os.path.isfile(local):
+        return False
+    dest = remote_dir.rstrip("/") + "/"
+    cmd = ["rsync", "-a", "-e",
+           "ssh -o IdentitiesOnly=yes -o StrictHostKeyChecking=accept-new",
+           local, dest]
+    try:
+        subprocess.check_call(cmd, stdout=subprocess.DEVNULL,
+                              stderr=subprocess.DEVNULL)
+        print(f"[hf_sync] staged {os.path.basename(local)} -> {dest}", flush=True)
+        return True
+    except Exception as exc:
+        print(f"[hf_sync] rsync to login failed: {exc}", flush=True)
+        return False
+
+
+def _upload_file(api, local, repo, dest, token):
+    if api is None:
+        print(f"[hf_sync] skip upload {dest}: huggingface_hub missing", flush=True)
+        return False
+    api.upload_file(
+        path_or_fileobj=local,
+        path_in_repo=dest,
+        repo_id=repo,
+        repo_type="model",
+        token=token,
+        commit_message=f"sync {dest}",
+    )
+    print(f"[hf_sync] uploaded {dest} ({os.path.getsize(local)} bytes)", flush=True)
+    return True
+
+
+def sync_once(args, state):
+    workdir = os.path.abspath(args.workdir)
+    scratch_log = args.scratch_log
+    if args.log and os.path.isfile(args.log) and scratch_log:
+        os.makedirs(os.path.dirname(scratch_log), exist_ok=True)
+        subprocess.call(["rsync", "-a", args.log, scratch_log])
+
+    hb_src = os.path.join(workdir, "heartbeat.json")
+    if os.path.isfile(hb_src) and scratch_log:
+        hb_scratch = os.path.join(os.path.dirname(scratch_log), "w12_inst_latent_heartbeat.json")
+        subprocess.call(["rsync", "-a", hb_src, hb_scratch])
+
+    token = (os.environ.get("HF_TOKEN")
+             or os.environ.get("HUGGING_FACE_HUB_TOKEN")
+             or "")
+    if args.token_file and os.path.isfile(args.token_file):
+        with open(args.token_file) as f:
+            token = f.read().strip() or token
+    if not token:
+        print("[hf_sync] no HF token; scratch copy only", flush=True)
+        return state
+
+    api = _hf_api(token)
+    prefix = args.prefix.strip("/")
+    repo = args.repo
+
+    if os.path.isfile(hb_src):
+        mtime = os.path.getmtime(hb_src)
+        if mtime > float(state.get("heartbeat_mtime", 0)):
+            if _upload_file(api, hb_src, repo, f"{prefix}/heartbeat.json", token):
+                state["heartbeat_mtime"] = mtime
+            _rsync_to_login(hb_src, args.login_stage)
+
+    if args.log and os.path.isfile(args.log):
+        mtime = os.path.getmtime(args.log)
+        size = os.path.getsize(args.log)
+        last = state.get("log", {})
+        due = (time.time() - float(last.get("uploaded_at", 0))) >= 300
+        if due and (mtime > float(last.get("mtime", 0)) or size != last.get("size", -1)):
+            if _upload_file(api, args.log, repo, f"{prefix}/train.log", token):
+                state["log"] = {"mtime": mtime, "size": size,
+                                "uploaded_at": time.time()}
+
+    ready = _load_json(os.path.join(workdir, "ckpt_ready.json"))
+    ckpt_dir, ckpt_step = _latest_ckpt_dir(workdir)
+    if ckpt_dir and ckpt_step > int(state.get("latest_step", -1)):
+        tar_path = os.path.join(args.tar_dir, "latest.tar")
+        print(f"[hf_sync] tarring {ckpt_dir} -> {tar_path}", flush=True)
+        _tar(ckpt_dir, tar_path)
+        uploaded = _upload_file(api, tar_path, repo, f"{prefix}/latest.tar", token)
+        staged = _rsync_to_login(tar_path, args.login_stage)
+        if uploaded or staged:
+            state["latest_step"] = ckpt_step
+        meta = {
+            "step": ckpt_step,
+            "event": ready.get("event", "periodic"),
+            "stage": ready.get("stage"),
+            "host": socket.gethostname(),
+            "job_id": os.environ.get("SLURM_JOB_ID", ""),
+            "src": ckpt_dir,
+        }
+        meta_path = os.path.join(args.tar_dir, "latest.json")
+        _write_json(meta_path, meta)
+        _upload_file(api, meta_path, repo, f"{prefix}/latest.json", token)
+        _rsync_to_login(meta_path, args.login_stage)
+        if uploaded:
+            state["latest_step"] = ckpt_step
+        keep = bool(ready.get("keep_snapshot")) or (ckpt_step > 0 and ckpt_step % 50000 == 0)
+        if keep and ckpt_step not in set(state.get("snapshots", [])):
+            dest = f"{prefix}/steps/step_{ckpt_step}.tar"
+            snap = os.path.join(args.tar_dir, f"step_{ckpt_step}.tar")
+            if snap != tar_path:
+                subprocess.call(["/bin/cp", "-f", tar_path, snap])
+            if _upload_file(api, tar_path, repo, dest, token):
+                state.setdefault("snapshots", []).append(ckpt_step)
+            _rsync_to_login(snap if os.path.isfile(snap) else tar_path, args.login_stage)
+
+    stage_dir = os.path.join(workdir, "stage_ckpts")
+    uploaded_stages = set(state.get("stages", []))
+    for path in _stage_ckpt_dirs(stage_dir):
+        name = os.path.basename(path)
+        if name in uploaded_stages:
+            continue
+        tar_path = os.path.join(args.tar_dir, f"{name}.tar")
+        print(f"[hf_sync] tarring stage {path}", flush=True)
+        _tar(path, tar_path)
+        if _upload_file(api, tar_path, repo, f"{prefix}/stages/{name}.tar", token):
+            uploaded_stages.add(name)
+        _rsync_to_login(tar_path, args.login_stage)
+    state["stages"] = sorted(uploaded_stages)
+    return state
+
+
+def main():
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--workdir", required=True)
+    ap.add_argument("--log", default="")
+    ap.add_argument("--scratch-log", default="")
+    ap.add_argument("--repo", default="Avra98/Sudoku_superposition")
+    ap.add_argument("--prefix", default="runs/w12_inst_latent")
+    ap.add_argument("--token-file", default="")
+    ap.add_argument("--state", default="")
+    ap.add_argument("--tar-dir", default="")
+    ap.add_argument("--interval", type=int, default=60)
+    ap.add_argument("--login-stage", default="gandalf.berkeley.edu:/tmp/sudoku_hf_uploads")
+    ap.add_argument("--once", action="store_true")
+    args = ap.parse_args()
+    args.tar_dir = args.tar_dir or os.path.join(os.path.abspath(args.workdir), "_hf_tars")
+    os.makedirs(args.tar_dir, exist_ok=True)
+    state_path = args.state or os.path.join(args.tar_dir, "hf_sync_state.json")
+    state = _load_json(state_path)
+    print(f"[hf_sync] host={socket.gethostname()} workdir={args.workdir}",
+          flush=True)
+    while True:
+        try:
+            state = sync_once(args, state)
+            _write_json(state_path, state)
+        except Exception:
+            traceback.print_exc()
+            print("[hf_sync] cycle failed; will retry", flush=True)
+        if args.once:
+            break
+        time.sleep(max(5, args.interval))
+    return 0
+
+
+if __name__ == "__main__":
+    sys.exit(main())

code/wavecurriculum_run/sbatch_instance_latent.sh

@@ -0,0 +1,191 @@
+#!/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
+echo "[$(date)] /tmp before cleanup:"
+df -h /tmp | tail -1
+# Drop leftover run dirs from dead jobs; keep env + instance/mask caches.
+rm -rf /tmp/sudoku_wave_runs
+mkdir -p "${LOCAL_LOG}"
+echo "[$(date)] /tmp after cleanup:"
+df -h /tmp | tail -1
+
+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
+# 12 stages x 40k is 480k of the 800k budget, so the hard cap can no longer
+# consume the whole run the way 80k did.
+export SUDOKU_PATIENCE=40000
+export SUDOKU_MIN_STAGE_STEPS=8000
+# Stage t -> t+1 requires val_mass >= 0.85: the probability mass the value head
+# puts on stage t's candidate set, measured teacher-forced. Chance is ~0.41 on
+# stage-0 sets. This replaces the old in-set rate, which was conditioned on an
+# exact positional location match and so was estimated from ~4% of steps.
+export SUDOKU_PROMOTE_ACC=0.85
+# Both extra gates start OFF: this run is the one that tells us the reachable
+# range of loc_wave and val_excess. Turn them on once we can read the scale.
+export SUDOKU_PROMOTE_LOC_WAVE=0.0
+export SUDOKU_PROMOTE_VAL_EXCESS=0.0
+export SUDOKU_MAX_STEPS="${SUDOKU_MAX_STEPS:-800000}"
+export SUDOKU_EVAL_EVERY=2000
+export SUDOKU_SAVE_EVERY=10000
+export SUDOKU_CKPT_KEEP=2
+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}"
+# Leave a little H200 headroom so eval + sidecar tar do not OOM the node.
+export XLA_PYTHON_CLIENT_MEM_FRACTION=0.85
+
+HF_TOKEN_FILE="${HF_TOKEN_FILE:-/scratch/users/gatmiry/.hf_token}"
+if [ -z "${HF_TOKEN:-}" ] && [ -s "${HF_TOKEN_FILE}" ]; then
+  HF_TOKEN="$(cat "${HF_TOKEN_FILE}")"
+  export HF_TOKEN
+fi
+export HUGGING_FACE_HUB_TOKEN="${HF_TOKEN:-}"
+HF_PKGS="${HF_PKGS:-/scratch/users/gatmiry/hf_pkgs}"
+TRAIN_LOG="${RUN_DIR}/logs/w12_inst_latent.log"
+SYNC_LOG="${RUN_DIR}/logs/hf_sync_w12_inst_latent.log"
+
+# Need room for 2 rolling ckpts (~1 GB) plus a tar in flight.
+tmp_avail_kb=$(df -Pk /tmp | awk 'NR==2{print $4}')
+echo "[$(date)] /tmp avail ${tmp_avail_kb} KB"
+if [ "${tmp_avail_kb}" -lt 3000000 ]; then
+  echo "ERROR: /tmp has less than 3G free; refusing to start" >&2
+  df -h /tmp
+  du -sh /tmp/* 2>/dev/null | sort -h | tail -20 >&2 || true
+  exit 1
+fi
+
+cd "${RUN_DIR}"
+# huggingface_hub lives on scratch (feanor /tmp is too full to pip-install).
+PYTHONPATH="${HF_PKGS}${PYTHONPATH:+:${PYTHONPATH}}" ${PY} -u "${RUN_DIR}/hf_sync.py" \
+  --workdir "${LOCAL_LOG}" \
+  --log "${TRAIN_LOG}" \
+  --scratch-log "${TRAIN_LOG}" \
+  --repo Avra98/Sudoku_superposition \
+  --prefix runs/w12_inst_latent \
+  --token-file "${HF_TOKEN_FILE}" \
+  --interval 60 \
+  > "${SYNC_LOG}" 2>&1 &
+SYNC_PID=$!
+
+echo "[$(date)] starting w12_inst_latent from scratch"
+echo "  K=12 recurrent=1 bt=0 aux=0 instance_dir=${INST_DIR}"
+echo "  workdir=${LOCAL_LOG}  log=${TRAIN_LOG}"
+echo "  hf=Avra98/Sudoku_superposition/runs/w12_inst_latent"
+# Log goes to scratch immediately so a /tmp-full crash still leaves a traceback.
+CUDA_VISIBLE_DEVICES=0 ${PY} -u -m train.main \
+  --workdir="${LOCAL_LOG}" --exp_name="w12_inst_latent" \
+  > "${TRAIN_LOG}" 2>&1
+EC=$?
+kill ${SYNC_PID} 2>/dev/null || true
+wait ${SYNC_PID} 2>/dev/null || true
+PYTHONPATH="${HF_PKGS}${PYTHONPATH:+:${PYTHONPATH}}" ${PY} -u "${RUN_DIR}/hf_sync.py" --once \
+  --workdir "${LOCAL_LOG}" \
+  --log "${TRAIN_LOG}" \
+  --scratch-log "${TRAIN_LOG}" \
+  --repo Avra98/Sudoku_superposition \
+  --prefix runs/w12_inst_latent \
+  --token-file "${HF_TOKEN_FILE}" \
+  >> "${SYNC_LOG}" 2>&1 || true
+echo "[$(date)] w12_inst_latent exit ${EC}"
+exit ${EC}

code/wavecurriculum_run/sbatch_multistage_super.sh

@@ -0,0 +1,188 @@
+#!/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_multi
+#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
+
+# Multi-stage superposition run, warm-started from the stage-1 checkpoint of
+# w12_s1_super and walking stages 2..12. Runs on a SECOND GPU; the stage-1 job
+# keeps going untouched.
+#
+# Every stage repeats the stage-1 recipe: that stage's own (puzzle, instance)
+# pairs are enumerated and walked in shuffled passes, so each instance of each
+# puzzle is seen at least 5 times before the stage can be left. The pool is
+# rebuilt on promotion because each stage has its own instances.
+#
+# Promotion needs BOTH superposition gates to clear, no plateau and no patience:
+#   mass*   >= SUDOKU_PROMOTE_ACC     mass sits inside the candidate set
+#   spread* >= SUDOKU_PROMOTE_SPREAD  and is spread evenly over it
+# plus SUDOKU_MIN_STAGE_STEPS, which is set to cover 5 passes at every stage
+# (stage 2 is the largest at 46,555 steps for 100k puzzles).
+#
+# IMPORTANT: this script must never wipe /tmp/sudoku_wave_runs, because the
+# stage-1 job's live checkpoints are in there.
+
+set -u
+hostname
+nvidia-smi -L
+echo "[$(date)] w12_multistage"
+
+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_multistage
+WARM_SRC=/tmp/sudoku_wave_runs/w12_s1_super
+WARM_DIR=/tmp/sudoku_warmstart_s1
+LOGIN_STAGE=gandalf.berkeley.edu:/tmp/sudoku_hf_multistage
+
+mkdir -p "${RUN_DIR}/logs" "${LOCAL_LOG}" "${WARM_DIR}"
+echo "[$(date)] /tmp free:"
+df -h /tmp | tail -1
+
+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 "ERROR: ${ENV_LOCAL} missing; the stage-1 job should have created it" >&2
+  exit 1
+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())"
+
+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
+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
+
+# ---- Warm start: snapshot the stage-1 job's newest checkpoint ----
+# Copied rather than read in place, so a concurrent write by the stage-1 job
+# cannot be observed half-finished, and so this run never writes into its dir.
+NEWEST=$(ls -1d ${WARM_SRC}/checkpoint_* 2>/dev/null \
+         | grep -v 'orbax-checkpoint-tmp' \
+         | sed 's/.*checkpoint_//' | sort -n | tail -1)
+if [ -z "${NEWEST}" ]; then
+  echo "ERROR: no checkpoint_* found in ${WARM_SRC}" >&2
+  ls -la "${WARM_SRC}" >&2 || true
+  exit 1
+fi
+echo "[$(date)] warm start from ${WARM_SRC}/checkpoint_${NEWEST}"
+rm -rf "${WARM_DIR:?}"/checkpoint_*
+cp -r "${WARM_SRC}/checkpoint_${NEWEST}" "${WARM_DIR}/"
+du -sh "${WARM_DIR}/checkpoint_${NEWEST}"
+
+export SUDOKU_RESUME=1
+export SUDOKU_START_STAGE=2
+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
+# Same sampling as stage 1, rebuilt per stage: every (puzzle, instance) pair of
+# the current stage, walked in shuffled passes, 5 sightings each.
+export SUDOKU_INSTANCE_EPOCHS=5
+# 100k puzzles keeps a stage at ~47k steps (~1.2 h), so all 11 stages fit in
+# ~15 h. The full corpus would be ~22 h for a single stage.
+export SUDOKU_INSTANCE_PUZZLES="${SUDOKU_INSTANCE_PUZZLES:-100000}"
+# Covers 5 passes at every stage; stage 2 needs the most at 46,555 steps.
+export SUDOKU_MIN_STAGE_STEPS=50000
+# Calibrated on what stage 1 actually reached: mass* 0.9957, spread* 0.92.
+export SUDOKU_PROMOTE_ACC=0.95
+export SUDOKU_PROMOTE_SPREAD=0.88
+export SUDOKU_PROMOTE_LOC_WAVE=0.0
+export SUDOKU_PLATEAU_STEPS=0
+export SUDOKU_PATIENCE=0
+# Resume continues from the restored optimizer step (~270k), plus 11 stages at
+# >=50k each.
+export SUDOKU_MAX_STEPS="${SUDOKU_MAX_STEPS:-950000}"
+export SUDOKU_EVAL_EVERY=2000
+export SUDOKU_SAVE_EVERY=10000
+export SUDOKU_CKPT_KEEP=2
+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}"
+# Two jobs share the H200s, so leave more headroom than the single-job setting.
+export XLA_PYTHON_CLIENT_MEM_FRACTION=0.45
+
+HF_TOKEN_FILE="${HF_TOKEN_FILE:-/scratch/users/gatmiry/.hf_token}"
+if [ -z "${HF_TOKEN:-}" ] && [ -s "${HF_TOKEN_FILE}" ]; then
+  HF_TOKEN="$(cat "${HF_TOKEN_FILE}")"
+  export HF_TOKEN
+fi
+export HUGGING_FACE_HUB_TOKEN="${HF_TOKEN:-}"
+HF_PKGS="${HF_PKGS:-/scratch/users/gatmiry/hf_pkgs}"
+TRAIN_LOG="${RUN_DIR}/logs/w12_multistage.log"
+SYNC_LOG="${RUN_DIR}/logs/hf_sync_w12_multistage.log"
+
+tmp_avail_kb=$(df -Pk /tmp | awk 'NR==2{print $4}')
+echo "[$(date)] /tmp avail ${tmp_avail_kb} KB"
+if [ "${tmp_avail_kb}" -lt 3000000 ]; then
+  echo "ERROR: /tmp has less than 3G free; refusing to start" >&2
+  exit 1
+fi
+
+cd "${RUN_DIR}"
+PYTHONPATH="${HF_PKGS}${PYTHONPATH:+:${PYTHONPATH}}" ${PY} -u "${RUN_DIR}/hf_sync.py" \
+  --workdir "${LOCAL_LOG}" \
+  --log "${TRAIN_LOG}" \
+  --scratch-log "${TRAIN_LOG}" \
+  --repo Avra98/Sudoku_superposition \
+  --prefix runs/w12_multistage \
+  --login-stage "${LOGIN_STAGE}" \
+  --token-file "${HF_TOKEN_FILE}" \
+  --interval 60 \
+  > "${SYNC_LOG}" 2>&1 &
+SYNC_PID=$!
+
+echo "[$(date)] starting w12_multistage"
+echo "  warm start from stage-1 checkpoint ${NEWEST}, stages 2..12"
+echo "  ${SUDOKU_INSTANCE_PUZZLES} puzzles, 5 passes per stage, min ${SUDOKU_MIN_STAGE_STEPS} steps"
+echo "  promote when mass* >= ${SUDOKU_PROMOTE_ACC} and spread* >= ${SUDOKU_PROMOTE_SPREAD}"
+echo "  hf=Avra98/Sudoku_superposition/runs/w12_multistage"
+CUDA_VISIBLE_DEVICES=0 ${PY} -u -m train.main \
+  --workdir="${LOCAL_LOG}" --exp_name="w12_multistage" \
+  --ckpt_loc="${WARM_DIR}" \
+  > "${TRAIN_LOG}" 2>&1
+EC=$?
+kill ${SYNC_PID} 2>/dev/null || true
+wait ${SYNC_PID} 2>/dev/null || true
+PYTHONPATH="${HF_PKGS}${PYTHONPATH:+:${PYTHONPATH}}" ${PY} -u "${RUN_DIR}/hf_sync.py" --once \
+  --workdir "${LOCAL_LOG}" \
+  --log "${TRAIN_LOG}" \
+  --scratch-log "${TRAIN_LOG}" \
+  --repo Avra98/Sudoku_superposition \
+  --prefix runs/w12_multistage \
+  --login-stage "${LOGIN_STAGE}" \
+  --token-file "${HF_TOKEN_FILE}" \
+  >> "${SYNC_LOG}" 2>&1 || true
+echo "[$(date)] w12_multistage exit ${EC}"
+exit ${EC}

code/wavecurriculum_run/sbatch_stage1_super.sh

@@ -0,0 +1,187 @@
+#!/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=w1_super
+#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
+
+# Stage-1 superposition experiment. ONE stage, no curriculum, no promotion.
+#
+# The point: a cell's candidate set is only visible to the model as the spread
+# of digits it sees at that cell across repeated showings of the SAME puzzle.
+# Random sampling showed a puzzle roughly once per 28k steps, and only one of
+# its instances, so that spread was never presented. Here every (puzzle,
+# instance) pair is enumerated and walked in shuffled order 5 times, so a batch
+# holds one instance of each of 64 different puzzles (p1-i1, p2-i2, ...) and one
+# epoch ends when every pair has appeared 5 times.
+
+set -u
+hostname
+nvidia-smi -L
+echo "[$(date)] w12_s1_super"
+
+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_s1_super
+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
+LOGIN_STAGE=gandalf.berkeley.edu:/tmp/sudoku_hf_s1super
+
+mkdir -p "${RUN_DIR}/logs" /tmp/sudoku_wave_runs
+echo "[$(date)] /tmp before cleanup:"
+df -h /tmp | tail -1
+# Drop leftover run dirs from dead jobs; keep env + instance/mask caches.
+rm -rf /tmp/sudoku_wave_runs
+mkdir -p "${LOCAL_LOG}"
+echo "[$(date)] /tmp after cleanup:"
+df -h /tmp | tail -1
+
+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
+
+# ---- Stage 1 only: start == max, so the curriculum can never promote ----
+export SUDOKU_RESUME=0
+export SUDOKU_START_STAGE=1
+export SUDOKU_MAX_STAGE=1
+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
+# One epoch = every (puzzle, instance) pair of stage 1 seen at least 5 times.
+# 1,804,462 puzzles x 4-9 stored instances = 11,336,803 pairs.
+export SUDOKU_INSTANCE_EPOCHS=5
+export SUDOKU_INSTANCE_PUZZLES="${SUDOKU_INSTANCE_PUZZLES:-0}"
+# steps = 5 * 11,336,803 / 64 = 885,688, rounded up so the epoch completes.
+export SUDOKU_MAX_STEPS="${SUDOKU_MAX_STEPS:-886000}"
+export SUDOKU_EVAL_EVERY=2000
+export SUDOKU_SAVE_EVERY=10000
+export SUDOKU_CKPT_KEEP=2
+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}"
+# Leave a little H200 headroom so eval + sidecar tar do not OOM the node.
+export XLA_PYTHON_CLIENT_MEM_FRACTION=0.85
+
+HF_TOKEN_FILE="${HF_TOKEN_FILE:-/scratch/users/gatmiry/.hf_token}"
+if [ -z "${HF_TOKEN:-}" ] && [ -s "${HF_TOKEN_FILE}" ]; then
+  HF_TOKEN="$(cat "${HF_TOKEN_FILE}")"
+  export HF_TOKEN
+fi
+export HUGGING_FACE_HUB_TOKEN="${HF_TOKEN:-}"
+HF_PKGS="${HF_PKGS:-/scratch/users/gatmiry/hf_pkgs}"
+TRAIN_LOG="${RUN_DIR}/logs/w12_s1_super.log"
+SYNC_LOG="${RUN_DIR}/logs/hf_sync_w12_s1_super.log"
+
+# Need room for 2 rolling ckpts (~1 GB) plus a tar in flight.
+tmp_avail_kb=$(df -Pk /tmp | awk 'NR==2{print $4}')
+echo "[$(date)] /tmp avail ${tmp_avail_kb} KB"
+if [ "${tmp_avail_kb}" -lt 3000000 ]; then
+  echo "ERROR: /tmp has less than 3G free; refusing to start" >&2
+  df -h /tmp
+  du -sh /tmp/* 2>/dev/null | sort -h | tail -20 >&2 || true
+  exit 1
+fi
+
+cd "${RUN_DIR}"
+PYTHONPATH="${HF_PKGS}${PYTHONPATH:+:${PYTHONPATH}}" ${PY} -u "${RUN_DIR}/hf_sync.py" \
+  --workdir "${LOCAL_LOG}" \
+  --log "${TRAIN_LOG}" \
+  --scratch-log "${TRAIN_LOG}" \
+  --repo Avra98/Sudoku_superposition \
+  --prefix runs/w12_s1_super \
+  --login-stage "${LOGIN_STAGE}" \
+  --token-file "${HF_TOKEN_FILE}" \
+  --interval 60 \
+  > "${SYNC_LOG}" 2>&1 &
+SYNC_PID=$!
+
+echo "[$(date)] starting w12_s1_super from scratch"
+echo "  stage pinned to 1 (start=max=1), 1 latent pass, aux=0"
+echo "  1 epoch = every (puzzle, instance) pair seen 5x"
+echo "  workdir=${LOCAL_LOG}  log=${TRAIN_LOG}"
+echo "  hf=Avra98/Sudoku_superposition/runs/w12_s1_super"
+CUDA_VISIBLE_DEVICES=0 ${PY} -u -m train.main \
+  --workdir="${LOCAL_LOG}" --exp_name="w12_s1_super" \
+  > "${TRAIN_LOG}" 2>&1
+EC=$?
+kill ${SYNC_PID} 2>/dev/null || true
+wait ${SYNC_PID} 2>/dev/null || true
+PYTHONPATH="${HF_PKGS}${PYTHONPATH:+:${PYTHONPATH}}" ${PY} -u "${RUN_DIR}/hf_sync.py" --once \
+  --workdir "${LOCAL_LOG}" \
+  --log "${TRAIN_LOG}" \
+  --scratch-log "${TRAIN_LOG}" \
+  --repo Avra98/Sudoku_superposition \
+  --prefix runs/w12_s1_super \
+  --login-stage "${LOGIN_STAGE}" \
+  --token-file "${HF_TOKEN_FILE}" \
+  >> "${SYNC_LOG}" 2>&1 || true
+echo "[$(date)] w12_s1_super exit ${EC}"
+exit ${EC}

code/wavecurriculum_run/sbatch_wave12.sh

@@ -0,0 +1,203 @@
+#!/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=wave12
+#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
+
+# Wave-depth curriculum, 12 stages.
+#
+# The curriculum axis is PROPAGATION DEPTH, not puzzle difficulty. Each puzzle's
+# ~25 solver waves are subsampled to 12 evenly spaced candidate snapshots
+# (staged_candidate_gen.py --stages 12, last snapshot = the unique solution).
+# Curriculum stage t means: t recurrence passes, latent slots 1..t active, wave
+# snapshots 1..t supervised. Promotion adds exactly one wave block.
+#
+# Difficulty is NOT gated: every puzzle is available from step 0, drawn
+# uniformly from the corpus. The old level<=stage+2 filter is gone -- the
+# difficulty tag has only 6 values (so it cannot express a 12-step ladder), is
+# uncorrelated with puzzle size (r=-0.003 vs empty cells), and explains only
+# ~19% of the variance in solver round count.
+#
+# Promotion is plateau-driven: a stage graduates when the candidate-set accuracy
+# at its deepest slot stops improving, scored on cells that CHANGED from the
+# previous snapshot. Unchanged cells are copies of slot j-1 and stay correct for
+# a head that learned nothing, so they cannot signal acquisition. The accuracy
+# threshold is a fast path and the patience cap prevents a stuck stage stalling.
+#
+# Submit one job per arm:
+#   sbatch --job-name=w12_flat      --export=ALL,ARM=flat      sbatch_wave12.sh
+#   sbatch --job-name=w12_datacur   --export=ALL,ARM=datacur   sbatch_wave12.sh
+#   sbatch --job-name=w12_latent    --export=ALL,ARM=latent    sbatch_wave12.sh
+#   sbatch --job-name=w12_latent_bt --export=ALL,ARM=latent_bt sbatch_wave12.sh
+#
+# ARM:
+#   flat       K=0, no curriculum at all: pure control.
+#   datacur    K=0 with the 12-stage round-count DATA curriculum. Stage t admits
+#              only puzzles in the first t round-count bins, so the propagation
+#              ladder is expressed in the puzzle POOL rather than in latent
+#              supervision. This is the no-latent curriculum arm: same axis as
+#              the latent arms, no latent tokens. Promotion is gated on accuracy
+#              over the newest round-bin.
+#   latent     K=12 latent chain, 12-stage wave supervision, plateau promotion.
+#   latent_bt  as latent, plus adaptive backtracking keyed on DEPTH: a repair
+#              replays num_passes=t when snapshot t's accuracy regresses below
+#              its graduation value.
+#
+# From scratch is forced, not chosen: pos_embeddings is exactly (3*81+K, emb_dim),
+# so no 6-slot checkpoint can restore into a 12-slot model.
+#
+# The 12-stage masks live on feanor's node-local /tmp (sbatch_gen_s12_masks.sh)
+# because the 20 G scratch quota has ~2 G free and the train masks are 3.5 G.
+
+set -u
+hostname
+nvidia-smi -L
+echo "[$(date)] ARM=${ARM:-unset}"
+
+ARM="${ARM:?set ARM=flat|datacur|latent|latent_bt via --export=ALL,ARM=...}"
+NAME="w12_${ARM}"
+
+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
+LOCAL_LOG=/tmp/sudoku_wave_runs/${NAME}
+TARBALL_GANDALF=/tmp/logicpuzzles_env.tar.gz
+TARBALL_LOCAL=/tmp/logicpuzzles_env_${SLURM_JOB_ID}.tar.gz
+
+mkdir -p "${RUN_DIR}/logs" "${LOCAL_LOG}" /tmp/sudoku_wave_runs
+
+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_OPTS="-o IdentitiesOnly=yes -o StrictHostKeyChecking=accept-new"
+  [ -f "${HOME}/.ssh/id_ed25519_berkeley" ] && SCP_OPTS="${SCP_OPTS} -i ${HOME}/.ssh/id_ed25519_berkeley"
+  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())"
+
+# ---- Shared recipe: identical across arms ----
+export SUDOKU_RESUME=0
+export SUDOKU_START_STAGE=1
+export SUDOKU_MAX_STAGE="${SUDOKU_MAX_STAGE:-12}"
+# Plateau promotion. A stage graduates when its frontier-depth accuracy has not
+# gained 0.005 for 20k steps; PATIENCE is only a hard ceiling. This replaces the
+# old fixed 15k timer, under which every promotion fired on patience at 0.32-0.46
+# accuracy and the 0.70 threshold was dead code.
+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.90
+# Matches the flat reference run (which reached level-3 acc 0.945 at ~800k), so
+# the arms are comparable to it at equal steps rather than stopping at 250k.
+export SUDOKU_MAX_STEPS="${SUDOKU_MAX_STEPS:-800000}"
+# Difficulty tag selects nothing: uniform over the corpus (~68% level 3).
+export SUDOKU_LEVEL_BALANCED=0
+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="../datasets/train_sudoku_puzzles.npy"
+export SUDOKU_TEST_PATH="../datasets/test_sudoku_puzzles.npy"
+export XLA_PYTHON_CLIENT_MEM_FRACTION=0.9
+
+# ---- Per-arm knobs ----
+export SUDOKU_BACKTRACK=0
+case "${ARM}" in
+  flat)
+    # No latents: no candidate head, no depth ladder. Masks left unset.
+    export SUDOKU_LATENT_SLOTS=0
+    export SUDOKU_RECURRENT=0
+    export SUDOKU_AUX_WEIGHT=0.0
+    export SUDOKU_TRAIN_CAND=""
+    export SUDOKU_TEST_CAND=""
+    ;;
+  datacur)
+    # No latents, but a real 12-stage curriculum over the puzzle pool, ordered
+    # by solver round count. Candidate masks stay unset (nothing to supervise);
+    # only the meta files are needed, for the round counts.
+    export SUDOKU_LATENT_SLOTS=0
+    export SUDOKU_RECURRENT=0
+    export SUDOKU_AUX_WEIGHT=0.0
+    export SUDOKU_TRAIN_CAND=""
+    export SUDOKU_TEST_CAND=""
+    export SUDOKU_DATA_CURRICULUM=rounds
+    export SUDOKU_TRAIN_META="${CAND_DIR}/train_meta.npy"
+    export SUDOKU_TEST_META="${CAND_DIR}/test_meta.npy"
+    for f in "${SUDOKU_TRAIN_META}" "${SUDOKU_TEST_META}"; do
+      [ -s "${f}" ] || { echo "missing meta ${f}; run sbatch_gen_s12_masks.sh on this node" >&2; exit 1; }
+    done
+    ;;
+  latent|latent_bt)
+    export SUDOKU_LATENT_SLOTS=12
+    export SUDOKU_RECURRENT=1
+    export SUDOKU_AUX_WEIGHT=1.0
+    export SUDOKU_CAND_SLOT_MODE=depth      # k = num_passes, slot j -> snapshot j
+    export SUDOKU_PASSES_PER_STAGE=1        # stage t -> depth t, 12 stages -> 12 slots
+    export SUDOKU_CAND_DELTA_BG=0.25        # down-weight cells copied from slot j-1
+    export SUDOKU_TRAIN_CAND="${CAND_DIR}/train_cand_masks.npy"
+    export SUDOKU_TEST_CAND="${CAND_DIR}/test_cand_masks.npy"
+    for f in "${SUDOKU_TRAIN_CAND}" "${SUDOKU_TEST_CAND}"; do
+      [ -s "${f}" ] || { echo "missing masks ${f}; run sbatch_gen_s12_masks.sh on this node" >&2; exit 1; }
+    done
+    ;;
+  *) echo "unknown ARM '${ARM}'" >&2; exit 1 ;;
+esac
+
+if [ "${ARM}" = "latent_bt" ]; then
+  # Adaptive replay keyed on depth: deficits are measured on snapshot t's
+  # candidate-set accuracy, and a repair trains at num_passes=t on the same
+  # full-corpus batch distribution the frontier uses.
+  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
+fi
+
+cd "${RUN_DIR}"
+# Log-only periodic sync: checkpoints stay on /tmp (20 G scratch quota).
+(
+  while true; do sleep 900
+    rsync -a "${LOCAL_LOG}.log" "${RUN_DIR}/logs/${NAME}.log" 2>/dev/null || true
+  done
+) &
+SYNC_PID=$!
+
+echo "[$(date)] starting ${NAME} from scratch"
+echo "  K=${SUDOKU_LATENT_SLOTS} recurrent=${SUDOKU_RECURRENT} bt=${SUDOKU_BACKTRACK}"
+echo "  max_stage=${SUDOKU_MAX_STAGE} plateau=${SUDOKU_PLATEAU_STEPS} patience=${SUDOKU_PATIENCE}"
+echo "  slot_mode=${SUDOKU_CAND_SLOT_MODE:-n/a} pps=${SUDOKU_PASSES_PER_STAGE:-n/a} delta_bg=${SUDOKU_CAND_DELTA_BG:-n/a}"
+CUDA_VISIBLE_DEVICES=0 ${PY} -u -m train.main \
+  --workdir="${LOCAL_LOG}" --exp_name="${NAME}" \
+  > "${LOCAL_LOG}.log" 2>&1
+EC=$?
+kill ${SYNC_PID} 2>/dev/null || true
+rsync -a "${LOCAL_LOG}.log" "${RUN_DIR}/logs/${NAME}.log" 2>/dev/null || true
+echo "[$(date)] ${NAME} exit ${EC}"
+exit ${EC}

code/wavecurriculum_run/show_val_example.py

@@ -0,0 +1,63 @@
+"""Worked example: one cell with candidate set S = {1,2,4}, four model behaviours.
+
+Shows that excess = log(1/mass) + KL(Uniform(S) || p/mass) separates "support
+leaks outside S" from "support inside S is not uniform".
+"""
+import sys
+import os
+import types
+
+import numpy as np
+
+for name in ("jax", "jax.numpy", "flax", "flax.training",
+             "flax.training.common_utils", "train.model"):
+    sys.modules.setdefault(name, types.ModuleType(name))
+sys.modules["jax"].numpy = sys.modules["jax.numpy"]
+sys.modules["flax"].training = sys.modules["flax.training"]
+sys.modules["flax.training"].common_utils = sys.modules[
+    "flax.training.common_utils"]
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from train.evaluater import _cand_value_stats
+
+CAND = [1, 2, 4]                      # digits
+BITS = int(sum(1 << (d - 1) for d in CAND))
+
+cases = [
+    ("A perfect superposition   ", {1: 1 / 3, 2: 1 / 3, 4: 1 / 3}),
+    ("B uniform over all 9      ", {d: 1 / 9 for d in range(1, 10)}),
+    ("C in S but collapsed on 4 ", {1: 0.05, 2: 0.05, 4: 0.90}),
+    ("D leaks AND collapsed     ", {1: 0.05, 2: 0.05, 4: 0.50,
+                                    7: 0.20, 8: 0.20}),
+    # A softmax never emits an exact zero, so keep a little mass on S here;
+    # with p(d) == 0 for every candidate the CE is +inf by definition.
+    ("E almost all mass outside S", {1: 0.005, 2: 0.005, 4: 0.005, 7: 0.985}),
+]
+
+hdr = (f"{'case':<28}{'mass':>7}{'out':>8}{'leak':>9}{'kl':>9}"
+       f"{'excess':>9}{'spread':>8}{'score':>8}")
+print(f"candidate set S = {{{','.join(map(str, CAND))}}}, |S| = {len(CAND)}, "
+      f"log|S| = {np.log(len(CAND)):.4f}")
+print(hdr)
+print("-" * len(hdr))
+for label, pmap in cases:
+    p = np.zeros(9)
+    for d, v in pmap.items():
+        p[d - 1] = v
+    assert abs(p.sum() - 1.0) < 1e-9, (label, p.sum())
+    logp = np.log(np.maximum(p, 1e-300))[None, :]
+    st = _cand_value_stats(logp, np.array([BITS]))
+    excess = st["ce"] - st["floor"]
+    mass = st["mass"]
+    leak = np.log(1.0 / max(mass, 1e-300))
+    kl = st["kl_multi"]
+    # the identity, recomputed independently of the code
+    assert abs(excess - (leak + kl)) < 1e-8, (label, excess, leak + kl)
+    print(f"{label:<28}{mass:>7.3f}{1 - mass:>8.3f}{leak:>9.3f}{kl:>9.3f}"
+          f"{excess:>9.3f}{st['spread']:>8.3f}{np.exp(-excess):>8.3f}")
+
+print()
+print("out    = val_acc          = 1 - mass  = P(digit outside S)")
+print("leak   = log(1/mass)      -> 0 when no support outside S")
+print("kl     = KL(Unif(S)||p/m) -> 0 when support inside S is uniform")
+print("excess = leak + kl        -> 0 iff BOTH; this is the promotion signal")
+print("score  = exp(-excess)     -> 1.0 iff both; compared to SUDOKU_PROMOTE_ACC")

code/wavecurriculum_run/test_instance_epochs.py

@@ -0,0 +1,83 @@
+"""Check the (puzzle, instance) epoch enumeration on the real stage-1 counts.
+
+Verifies the property the experiment rests on: after E epochs every instance of
+every puzzle has been used exactly E times, and a batch holds distinct puzzles
+(one instance of a given puzzle per batch).
+"""
+import numpy as np
+
+
+def build_pairs(starts, counts, stage, n_puz):
+    """Same expansion as SudokuDataset._build_instance_epoch_list."""
+    c = np.asarray(counts[:n_puz, stage]).astype(np.int64)
+    s = np.asarray(starts[:n_puz, stage]).astype(np.int64)
+    total = int(c.sum())
+    puzzle_ids = np.repeat(np.arange(n_puz, dtype=np.int64), c)
+    offsets = np.arange(total, dtype=np.int64) - np.repeat(np.cumsum(c) - c, c)
+    rows = np.repeat(s, c) + offsets
+    return puzzle_ids, rows, c
+
+
+def main():
+    d = "/tmp/sudoku_superposition"
+    starts = np.load(f"{d}/train_starts.npy")
+    counts = np.load(f"{d}/train_counts.npy")
+    stage, n_puz, epochs, bs = 0, 20000, 5, 64
+
+    pids, rows, c = build_pairs(starts, counts, stage, n_puz)
+    total = len(rows)
+    print(f"stage {stage + 1}: {n_puz} puzzles -> {total} pairs, "
+          f"{c.min()}-{c.max()} instances/puzzle (mean {c.mean():.2f})")
+
+    # 1. rows are exactly each puzzle's contiguous assignment block, no repeats
+    assert total == int(c.sum())
+    assert len(np.unique(rows)) == total, "duplicate assignment rows"
+    for p in (0, 1, 7, n_puz - 1):
+        want = np.arange(starts[p, stage], starts[p, stage] + counts[p, stage])
+        got = np.sort(rows[pids == p])
+        assert np.array_equal(got, want), (p, got, want)
+    print("rows match each puzzle's assignment block            OK")
+
+    # 2. every instance seen exactly `epochs` times
+    rng = np.random.RandomState(0)
+    seen = np.zeros(total, dtype=np.int32)
+    order = np.arange(total)
+    batch_dup = 0
+    n_batches = 0
+    for _ in range(epochs):
+        rng.shuffle(order)
+        seen[order] += 1
+        for b in range(0, total - bs, bs):
+            sl = pids[order[b:b + bs]]
+            batch_dup += bs - len(np.unique(sl))
+            n_batches += 1
+    assert seen.min() == seen.max() == epochs, (seen.min(), seen.max())
+    print(f"every instance seen exactly {epochs}x                     OK")
+
+    # 3. per-puzzle: all of its instances covered, each `epochs` times
+    per_puzzle = np.bincount(pids, weights=seen, minlength=n_puz)
+    assert np.array_equal(per_puzzle, epochs * c)
+    print("every puzzle: all instances x epochs                 OK")
+
+    # 4. batches hold distinct puzzles
+    print(f"same-puzzle collisions in a batch: {batch_dup} over "
+          f"{n_batches} batches ({batch_dup / (n_batches * bs) * 100:.3f}% "
+          f"of slots)")
+
+    # 5. averaged over the epochs, the target at a cell IS the candidate set
+    A = np.load(f"{d}/train_assignments.npy", mmap_mode="r")
+    p = 3
+    blk = np.asarray(A[starts[p, stage]:starts[p, stage] + counts[p, stage]])
+    multi = [c_ for c_ in range(81) if len(np.unique(blk[:, c_])) > 1]
+    print(f"\npuzzle {p}: {counts[p, stage]} instances, "
+          f"{len(multi)} cells whose digit varies across instances")
+    for cell in multi[:5]:
+        vals, cnt = np.unique(blk[:, cell], return_counts=True)
+        emp = ", ".join(f"{v}:{n}/{len(blk)}" for v, n in zip(vals, cnt))
+        print(f"  cell (r{cell // 9}, c{cell % 9}) digits {emp}")
+    print("\nthis empirical spread over instances is exactly what the model "
+          "must reproduce; it only sees it if all instances are visited")
+
+
+if __name__ == "__main__":
+    main()

code/wavecurriculum_run/test_new_metrics.py

@@ -0,0 +1,163 @@
+"""Numerical checks for the new location and value metrics.
+
+Run: python wavecurriculum_run/test_new_metrics.py
+"""
+import sys
+import os
+import types
+
+import numpy as np
+
+# The helpers under test are pure numpy, so stub the heavy imports at the top of
+# evaluater.py to keep this runnable on any node without jax installed.
+for name in ("jax", "jax.numpy", "flax", "flax.training",
+             "flax.training.common_utils", "train.model"):
+    sys.modules.setdefault(name, types.ModuleType(name))
+sys.modules["jax"].numpy = sys.modules["jax.numpy"]
+sys.modules["flax"].training = sys.modules["flax.training"]
+sys.modules["flax.training"].common_utils = sys.modules[
+    "flax.training.common_utils"]
+
+sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
+from train.evaluater import (_cand_value_stats, _wave_order_score, _lcs_len)
+
+fails = []
+
+
+def check(name, got, want, tol=1e-6):
+    ok = abs(float(got) - float(want)) <= tol
+    print(f"{'PASS' if ok else 'FAIL'}  {name}: got {got:.6f} want {want:.6f}")
+    if not ok:
+        fails.append(name)
+
+
+def logp_from_probs(p):
+    p = np.asarray(p, dtype=np.float64)
+    return np.log(np.maximum(p, 1e-300))
+
+
+BITS_123 = 0b000000111          # candidates {1,2,3}
+BITS_5 = 0b000010000            # candidate {5}
+
+# ---- val_excess == 0 exactly at uniform over the candidate set ----
+p = np.zeros((1, 9)); p[0, :3] = 1.0 / 3.0
+st = _cand_value_stats(logp_from_probs(p), np.array([BITS_123]))
+check("uniform-over-S excess", st["ce"] - st["floor"], 0.0)
+check("uniform-over-S mass", st["mass"], 1.0)
+check("uniform-over-S spread", st["spread"], 1.0)
+check("uniform-over-S floor==log3", st["floor"], np.log(3))
+
+# ---- collapse onto one candidate: excess > 0, spread -> 0 ----
+p = np.zeros((1, 9)); p[0, 0] = 0.98; p[0, 1] = 0.01; p[0, 2] = 0.01
+st = _cand_value_stats(logp_from_probs(p), np.array([BITS_123]))
+excess = st["ce"] - st["floor"]
+print(f"      collapsed excess = {excess:.4f} (want >> 0), "
+      f"spread = {st['spread']:.4f} (want ~0)")
+if not excess > 1.0:
+    fails.append("collapsed excess > 1")
+if not st["spread"] < 0.2:
+    fails.append("collapsed spread < 0.2")
+
+# ---- leakage outside the set lowers mass and raises excess ----
+p = np.zeros((1, 9)); p[0, :3] = 0.2; p[0, 5] = 0.4
+st = _cand_value_stats(logp_from_probs(p), np.array([BITS_123]))
+check("leaky mass", st["mass"], 0.6)
+if not (st["ce"] - st["floor"]) > 0.4:
+    fails.append("leaky excess")
+print(f"      leaky excess = {st['ce'] - st['floor']:.4f} (want > 0.4)")
+
+# ---- uniform over all 9 digits: excess == log(9/|S|) ----
+p = np.full((1, 9), 1.0 / 9.0)
+st = _cand_value_stats(logp_from_probs(p), np.array([BITS_123]))
+check("uniform-over-9 excess", st["ce"] - st["floor"], np.log(9) - np.log(3))
+check("uniform-over-9 mass", st["mass"], 3.0 / 9.0)
+check("uniform-over-9 spread", st["spread"], 1.0)
+
+# ---- singleton set is excluded from spread but counted elsewhere ----
+p = np.zeros((1, 9)); p[0, 4] = 1.0
+st = _cand_value_stats(logp_from_probs(p), np.array([BITS_5]))
+check("singleton excess", st["ce"] - st["floor"], 0.0)
+check("singleton spread_count", st["spread_count"], 0)
+check("singleton count", st["count"], 1)
+
+# ---- the decomposition identity, on many random distributions ----
+#   excess = log(1/mass) + KL(Uniform(S) || p/mass)
+# This is the whole claim: one number that is zero iff the support outside the
+# candidate set has vanished AND the support inside it is uniform.
+rng = np.random.default_rng(0)
+worst = 0.0
+for trial in range(2000):
+    k = int(rng.integers(2, 7))
+    cand = rng.choice(9, size=k, replace=False)
+    bits = int(sum(1 << int(d) for d in cand))
+    p = rng.dirichlet(np.full(9, float(rng.choice([0.15, 1.0, 5.0]))))
+    st = _cand_value_stats(logp_from_probs(p[None, :]), np.array([bits]))
+    excess = st["ce"] - st["floor"]
+    mass = st["mass"]
+    q = p[cand] / mass
+    kl = float(np.sum((1.0 / k) * np.log((1.0 / k) / q)))
+    worst = max(worst, abs(excess - (np.log(1.0 / mass) + kl)))
+    # both terms are non-negative, so excess is too
+    if excess < -1e-9 or kl < -1e-9:
+        fails.append("decomposition non-negativity")
+    # the code's own kl_multi must equal the independent one
+    if abs(st["kl_multi"] - kl) > 1e-8:
+        fails.append(f"kl_multi mismatch trial {trial}")
+        break
+print(f"{'PASS' if worst < 1e-8 else 'FAIL'}  decomposition identity over 2000 "
+      f"random cases: max |excess - (log(1/mass) + KL)| = {worst:.2e}")
+if worst >= 1e-8:
+    fails.append("decomposition identity")
+
+# ---- the two failure modes are separated by the split ----
+S = np.array([0b000010111])          # candidates {1,2,3,5}, |S| = 4
+# (a) perfectly uniform inside S, but only 50% of the mass is inside
+p = np.zeros((1, 9)); p[0, [0, 1, 2, 4]] = 0.125; p[0, [6, 7]] = 0.25
+st = _cand_value_stats(logp_from_probs(p), S)
+check("leak-only: kl term is 0", st["kl_multi"], 0.0)
+check("leak-only: excess == log(1/mass)", st["ce"] - st["floor"], np.log(2.0))
+# (b) all the mass inside S, but skewed across the candidates
+p = np.zeros((1, 9)); p[0, 0] = 0.7; p[0, 1] = 0.1; p[0, 2] = 0.1; p[0, 4] = 0.1
+st = _cand_value_stats(logp_from_probs(p), S)
+check("skew-only: mass is 1", st["mass"], 1.0)
+check("skew-only: excess == kl", st["ce"] - st["floor"], st["kl_multi"])
+print(f"      skew-only kl = {st['kl_multi']:.4f} (want > 0), "
+      f"spread = {st['spread']:.4f}")
+if not st["kl_multi"] > 0.2:
+    fails.append("skew-only kl > 0.2")
+
+# ---- singleton cells are excluded from the *_multi aggregates ----
+st = _cand_value_stats(logp_from_probs(np.full((1, 9), 1.0 / 9.0)),
+                       np.array([0b000010000]))
+check("singleton excluded from excess_multi", st["excess_multi"], 0.0)
+check("singleton excluded from mass_multi", st["mass_multi"], 0.0)
+check("singleton excluded from kl_multi", st["kl_multi"], 0.0)
+
+# ---- loc_wave: order up to ties within a wave ----
+truth = [(0, 0), (0, 1), (0, 2), (1, 0)]
+waves = np.array([0, 0, 1, 1])
+
+check("wave exact order", _wave_order_score(truth, truth, waves, 12), 4)
+# swapping within a wave must cost nothing
+check("wave swap inside wave",
+      _wave_order_score([(0, 1), (0, 0), (1, 0), (0, 2)], truth, waves, 12), 4)
+# pulling a later wave forward costs exactly those steps
+check("wave 1 pulled ahead of wave 0",
+      _wave_order_score([(0, 2), (0, 0), (0, 1), (1, 0)], truth, waves, 12), 3)
+# duplicates and non-cells score nothing and consume nothing
+check("wave duplicates",
+      _wave_order_score([(0, 0), (0, 0), (0, 0), (0, 0)], truth, waves, 12), 1)
+check("wave invalid cells",
+      _wave_order_score([(7, 7), (8, 8), (0, 0), (0, 1)], truth, waves, 12), 2)
+
+# ---- LCS: a reordering is not wiped out the way a positional match is ----
+check("lcs identical", _lcs_len(truth, truth), 4)
+check("lcs one swap",
+      _lcs_len([(0, 1), (0, 0), (0, 2), (1, 0)], truth), 3)
+check("lcs reversed", _lcs_len(truth[::-1], truth), 1)
+
+print()
+if fails:
+    print(f"{len(fails)} FAILURES: {fails}")
+    sys.exit(1)
+print("all metric checks passed")

code/wavecurriculum_run/train/data.py

@@ -0,0 +1,692 @@
+# 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()
+        self._build_instance_epoch_list()
+
+    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 _build_instance_epoch_list(self):
+        """Schedule for the pinned stage, so each puzzle is shown many times.
+
+        Default sampling draws a puzzle uniformly and then one of its
+        assignments at random, so over 8k steps at batch 64 a given puzzle is
+        seen at most once and most of its instances are never seen at all. The
+        superposition at a cell is only visible to the model as the spread of
+        digits it sees at that cell across repeated showings of the SAME puzzle,
+        so the schedule is built explicitly and walked in shuffled epochs.
+
+        Two sources of instances:
+
+        uniform (instance_uniform_draws=N)
+            Synthesize N instances per puzzle on the fly, each empty cell drawn
+            uniformly and independently from its candidate set. The per-cell
+            digit frequencies are then equal by construction, so the
+            cross-entropy optimum at that cell IS the uniform superposition.
+
+        pool (instance_epochs=E)
+            Walk the stored assignments E times each. Their per-cell frequencies
+            are whatever the coverage-driven generator produced (a |S|=2 cell is
+            typically 4:1), so CE converges to that skew, not to uniform.
+        """
+        self.inst_pairs = None
+        self.inst_uniform = 0
+        self.inst_pairs_stage = None
+        draws = int(getattr(self.config, "instance_uniform_draws", 0))
+        epochs = int(getattr(self.config, "instance_epochs", 0))
+        if not self.train or (draws <= 0 and epochs <= 0):
+            return
+        self._build_pairs_for_stage(self.instance_stage())
+
+    def _build_pairs_for_stage(self, stage):
+        """(Re)build the schedule for one stage.
+
+        Called again on promotion: each stage has its own instance pool, so the
+        pair list and the pass counter both restart when the stage advances.
+        """
+        draws = int(getattr(self.config, "instance_uniform_draws", 0))
+        epochs = int(getattr(self.config, "instance_epochs", 0))
+        n_puz = int(getattr(self.config, "instance_puzzles", 0))
+        n_all = len(self.train_puzzles)
+        n_puz = n_all if n_puz <= 0 else min(n_puz, n_all)
+
+        if draws > 0:
+            if self.cand_masks is None:
+                raise ValueError("instance_uniform_draws needs the candidate "
+                                 "masks (SUDOKU_TRAIN_CAND)")
+            self.inst_uniform = draws
+            widths = self._stage_widths(stage, min(n_puz, 2000))
+            pids = np.repeat(np.arange(n_puz, dtype=np.int32), draws)
+            # -1 = synthesize this instance instead of reading a stored row.
+            self.inst_pairs = (pids, np.full(len(pids), -1, dtype=np.int64))
+            print(f"[inst-uniform] stage {stage + 1}: {n_puz} puzzles x "
+                  f"{draws} uniform draws = {len(pids)} examples; mean |S| "
+                  f"{widths.mean():.2f}, max {widths.max()}; each candidate of "
+                  f"a cell is drawn ~{draws / widths.mean():.1f} times "
+                  f"(>=5 needs {int(5 * widths.max())} draws for the widest "
+                  f"cell)", flush=True)
+            self.inst_pairs_stage = stage
+            return
+
+        if self.instances is None:
+            raise ValueError("instance_epochs needs SUDOKU_INSTANCE_DIR")
+        counts = np.asarray(self.inst_counts[:n_puz, stage]).astype(np.int64)
+        starts = np.asarray(self.inst_starts[:n_puz, stage]).astype(np.int64)
+        total = int(counts.sum())
+        # Expand (puzzle -> its `count` consecutive assignment rows) without a
+        # Python loop: repeat the puzzle id, then add the within-puzzle offset.
+        puzzle_ids = np.repeat(np.arange(n_puz, dtype=np.int64), counts)
+        offsets = (np.arange(total, dtype=np.int64)
+                   - np.repeat(np.cumsum(counts) - counts, counts))
+        rows = np.repeat(starts, counts) + offsets
+        self.inst_pairs = (puzzle_ids.astype(np.int32), rows.astype(np.int64))
+        self.inst_epochs = epochs
+        self.inst_pairs_stage = stage
+        print(f"[inst] stage {stage + 1}: {n_puz} puzzles, {total} "
+              f"(puzzle, instance) pairs, {counts.min()}-{counts.max()} "
+              f"instances per puzzle (mean {counts.mean():.2f}); one epoch = "
+              f"every pair seen {epochs}x = {epochs * total} examples",
+              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 = int(getattr(self, "inst_stages", 0)) or int(self.num_stages)
+        if self.curriculum is None:
+            return S - 1
+        t = int(np.clip(self.curriculum.stage, 1, S))
+        return t - 1
+
+    def _stage_widths(self, stage, n_sample):
+        """|S| for every multi-candidate cell over the first n_sample puzzles."""
+        m = np.asarray(self.cand_masks[:n_sample, stage]).astype(np.int64)
+        bits = ((m[..., None] >> np.arange(9)) & 1).sum(-1)
+        return bits[bits >= 2]
+
+    def uniform_instance_values(self, idx, stage):
+        """Synthesize one assignment: each cell drawn uniformly from its set.
+
+        Returns (81,) digits indexed by cell = r*9+c. Cells the mask leaves
+        empty (mask 0) return 0, and the caller keeps the base sequence's digit
+        there, so the clue block is untouched.
+
+        The uniform draw over set bits is done by giving every set bit an iid
+        random key and taking the argmax: the max is equally likely to land on
+        any set bit, which is exactly a uniform choice, and it vectorizes over
+        all 81 cells at once.
+        """
+        m = np.asarray(self.cand_masks[idx, stage]).astype(np.int64)   # (81,)
+        bits = ((m[:, None] >> np.arange(9)) & 1).astype(np.float64)   # (81, 9)
+        keys = self.rng.random_sample((81, 9)) * bits
+        vals = (keys.argmax(1) + 1).astype(np.int8)
+        return np.where(m > 0, vals, 0)
+
+    def instance_values(self, idx, stage, inst_row=None):
+        """One assignment for (puzzle idx, stage).
+
+        inst_row pins an exact assignment row (epoch mode, so every instance is
+        visited a fixed number of times), -1 synthesizes a fresh uniform draw,
+        and None samples one of the stored rows at random.
+        """
+        if inst_row is not None and int(inst_row) < 0:
+            return self.uniform_instance_values(idx, stage)
+        if inst_row is not None:
+            return np.asarray(self.instances[int(inst_row)])
+        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, inst_row=None):
+        """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, inst_row=inst_row)
+        if vals is None:
+            return seq
+        seq = seq.copy()
+        cells = seq[0::3].astype(np.int64) * 9 + seq[1::3].astype(np.int64)
+        new = vals[cells].astype(seq.dtype)
+        # 0 means "this cell has no candidate mask"; keep the base digit there
+        # so the clue block survives untouched. Stored assignments carry the
+        # clue digit itself, so they overwrite with the same value either way.
+        seq[2::3] = np.where(new > 0, new, seq[2::3])
+        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, inst_row=None):
+        """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 or getattr(self, "inst_uniform", 0)) \
+                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(),
+                                       inst_row=inst_row)
+        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)
+
+        if getattr(self, "inst_pairs", None) is not None:
+            # Epoch mode: walk every (puzzle, instance) pair, reshuffled each
+            # epoch, so each instance of each puzzle is visited exactly once per
+            # epoch. Shuffling means a batch holds distinct puzzles, i.e. one
+            # instance of a given puzzle per batch rather than all of its
+            # instances side by side.
+            reps = max(int(getattr(self, "inst_epochs", 1)), 1)
+            npass = 0
+            while True:
+                stage = self.instance_stage()
+                if stage != self.inst_pairs_stage:
+                    # Promotion: this stage has its own instance pool, so the
+                    # pair list and the pass count both restart.
+                    self._build_pairs_for_stage(stage)
+                    npass = 0
+                puzzle_ids, rows = self.inst_pairs
+                total = len(rows)
+                order = np.arange(total)
+                self.rng.shuffle(order)
+                npass += 1
+                print(f"[inst] stage {stage + 1} pass {npass}: every "
+                      f"(puzzle, instance) pair seen {npass}x of {reps} "
+                      f"({total} pairs)", flush=True)
+                for i, t in enumerate(order):
+                    # Promotion can land mid-pass; checking periodically keeps
+                    # the targets on the current stage instead of finishing the
+                    # old pool first. 256 examples is 4 batches.
+                    if (i & 255) == 0 and self.instance_stage() != stage:
+                        break
+                    yield self.__getitem__(int(puzzle_ids[t]),
+                                           inst_row=int(rows[t]))
+
+        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/wavecurriculum_run/train/evaluater.py

@@ -0,0 +1,723 @@
+"""Evaluation related functions."""
+
+import os
+
+from flax.training import common_utils
+import jax
+from jax import numpy as jnp
+import numpy as np
+
+from train import model
+
+import pdb
+
+
+def _verbose_eval():
+    return os.environ.get("SUDOKU_VERBOSE_EVAL", "0") == "1"
+
+
+def _lcs_len(a, b):
+    """Length of the longest common subsequence of two cell sequences.
+
+    Measures order agreement while tolerating insertions, so a single early
+    deviation costs one step instead of desyncing every later comparison the
+    way a positional match does.
+    """
+    prev = [0] * (len(b) + 1)
+    for x in a:
+        cur = [0]
+        for k, y in enumerate(b):
+            cur.append(prev[k] + 1 if x == y else max(cur[k], prev[k + 1]))
+        prev = cur
+    return prev[-1]
+
+
+def _cand_value_stats(logp_digits, bits):
+    """Value-head statistics against one stage's candidate sets.
+
+    The instance target at a cell is drawn uniformly from that cell's candidate
+    set S, so the CE the trainer minimizes is
+        CE = -(1/|S|) sum_{d in S} log p(d)  >=  log|S|,
+    with equality iff p == Uniform(S); collapsing onto one candidate sends it to
+    infinity. So CE - log|S| == KL(Uniform(S) || p) is zero exactly at true
+    superposition and grows from leakage outside S or collapse inside it.
+
+    Args:
+        logp_digits: (n, 9) log-probabilities of digits 1..9 under the model.
+        bits: (n,) int bitmask; bit d-1 set iff digit d is a candidate.
+
+    Returns:
+        Dict of sums over the n cells plus the counts they divide by. `spread`
+        (normalized within-set entropy, 1.0 = uniform over S) is only defined
+        for |S| >= 2, hence its own count.
+    """
+    in_s = ((bits[:, None] >> np.arange(9)) & 1).astype(np.float64)  # (n, 9)
+    size = in_s.sum(axis=1)
+    p_d = np.exp(logp_digits)
+    mass = (in_s * p_d).sum(axis=1)
+    ce = -(in_s * logp_digits).sum(axis=1) / size
+    # excess splits exactly into the two things we want driven to zero:
+    #   excess = CE - log|S| = log(1/mass) + KL(Uniform(S) || p/mass)
+    #            \_________/   \_________/
+    #             leak outside   non-uniformity inside
+    # Both terms are >= 0, so excess == 0 iff all mass is on S AND spread
+    # uniformly over it. Derivation: writing p(d) = mass * q(d) for d in S,
+    # CE = -log(mass) - (1/|S|) sum log q(d) = -log(mass) + log|S| + KL(U||q).
+    excess = ce - np.log(size)
+    out = {
+        "ce": float(ce.sum()),
+        "floor": float(np.log(size).sum()),
+        "mass": float(mass.sum()),
+        "count": int(size.size),
+        "spread": 0.0,
+        "spread_count": 0,
+        # Same quantities restricted to genuinely superposed cells (|S| >= 2).
+        # Singleton cells are already determined at this stage, so there is
+        # nothing to be uniform about; as the curriculum deepens they come to
+        # dominate the pooled average and dilute the superposition signal.
+        "excess_multi": 0.0,
+        "mass_multi": 0.0,
+        "kl_multi": 0.0,
+    }
+    multi = size >= 2
+    if multi.any():
+        m = np.maximum(mass[multi], 1e-12)
+        q = in_s[multi] * p_d[multi] / m[:, None]
+        h = -(q * np.log(np.maximum(q, 1e-12))).sum(axis=1)
+        out["spread"] = float((h / np.log(size[multi])).sum())
+        out["spread_count"] = int(multi.sum())
+        out["excess_multi"] = float(excess[multi].sum())
+        out["mass_multi"] = float(mass[multi].sum())
+        # KL(Uniform(S) || q): the non-uniformity term of the split above.
+        # Stricter than 1 - spread, because it weights every candidate equally
+        # and so blows up when the model drops one candidate to ~0.
+        out["kl_multi"] = float((excess[multi] + np.log(m)).sum())
+    return out
+
+
+def _wave_order_score(emitted, truth, waves, max_wave):
+    """How many emitted cells came from the earliest still-unfilled wave.
+
+    A cell's wave is the propagation depth that determines it. Cells sharing a
+    wave are order-interchangeable, so the solver order in the data is only one
+    valid linearization; crediting any cell from the frontier wave measures
+    "is the model respecting propagation depth" without penalizing an arbitrary
+    tie-break. Cells that repeat or are not real empty cells score nothing.
+    """
+    left = dict(zip(truth, waves))
+    per_wave = np.bincount(np.asarray(waves, dtype=np.int64),
+                          minlength=max_wave + 1)
+    min_w, ok = 0, 0
+    for cell in emitted:
+        while min_w <= max_wave and per_wave[min_w] == 0:
+            min_w += 1
+        if min_w > max_wave:
+            break
+        cw = left.pop(cell, None)
+        if cw is None:
+            continue
+        per_wave[cw] -= 1
+        ok += int(cw == min_w)
+    return ok
+
+
+def valid_solution(output_seq):
+    """
+    This function checks if the puzzle is a valid solution by verifying if
+    each row, column and box has all the numbers from 1 to 9.
+
+    Args:
+        output_seq: a numpy array of shape (243,) containing the sequence of
+            output numbers
+
+    Returns:
+        int: 1 if correct solution, otherwise returns 0
+    """
+    # rows[i, j] keeps track if ith row has received (j + 1) number
+    rows = np.zeros((9, 9))
+    # cols[i, j] keeps track if ith column has received (j + 1) number
+    cols = np.zeros((9, 9))
+    # boxes[i, j] keeps track if ith box has received (j + 1) number
+    boxes = np.zeros((9, 9))
+
+    for j in range(81):
+        # The row and column are in the range (0, 8) and puzzle values are in (1, 9)
+        if int(output_seq[3 * j]) >= 9:
+            return False
+        if int(output_seq[3 * j + 1]) >= 9:
+            return False
+        if int(output_seq[3 * j + 2]) > 9:
+            return False
+        
+        row_num = int(output_seq[3 * j])
+        col_num = int(output_seq[3 * j + 1])
+        
+        # Mark the number in the row, column and box
+        rows[row_num, int(output_seq[3 * j + 2] - 1)] += 1
+        cols[col_num, int(output_seq[3 * j + 2] - 1)] += 1
+        boxes[
+            int(3 * (row_num // 3) + (col_num // 3)), int(output_seq[3 * j + 2] - 1)
+        ] += 1
+
+    if np.all(rows) and np.all(cols) and np.all(boxes):
+        return True
+    else:
+        return False
+
+
+def eval_step(state, batch, latent_vals, slot_pos, latent_active, config):
+    pred_logits, hidden, cand_logits = model.TransformerLMHeadModel(config).apply(
+        {"params": state.params}, batch, latent_values=latent_vals,
+        latent_positions=slot_pos, latent_active=latent_active,
+        )
+    return pred_logits, hidden, cand_logits
+
+
+def verify_sudoku_board(puzzle, row_num, col_num, num):
+    """
+    Args:
+        puzzle (np.array): The correct Sudoku puzzle.
+        row_num (int): The row number (0-8).
+        col_num (int): The column number (0-8).
+        num (int): The number predicted at the specified row and column.
+
+    Raises:
+        AssertionError: If the row_num * 9 + col_num >= 81 or if the number at the specified row and column is not equal to the given number.
+    """
+    if row_num * 9 + col_num >= 81: 
+        assert False
+  
+    assert puzzle[row_num * 9 + col_num] == num
+
+
+def get_eval_metrics(state, eval_data_iter, p_eval_step, config):
+    """This function computes given evaluation metrics (e.g, accuracy) in eval metrics for each batch and appends the metric in the list of eval_metrics.
+
+    Args: 
+        state: contains model parameters, optimizer, etc.
+        eval_data_iter: data iterator for evaluation dataset
+        p_eval_step: pmap function for forward pass of model for evaluation
+        config: general experiment config file
+
+    Returns: 
+        eval_metrics: contains list of evaluation metrics for each batch
+    """
+
+    eval_metrics = {
+        "acc": [],  # Unique-solution placement (parent leftover; not printed as val_acc)
+        "loc_acc": [],  # Location acc: model picks the ground-truth next cell (r,c)
+        # Permutation-tolerant location diagnostics. The upstream code scored
+        # only "acc" -- the digit at the cell the model chose -- and never
+        # compared (r,c) to the target order, because the solver order is one
+        # arbitrary linearization of a partial order: cells that become
+        # determined in the same propagation wave are interchangeable. loc_acc
+        # above therefore reads ~0.04 even when the model emits a valid
+        # permutation of the right cells, since one early deviation desyncs the
+        # rest of the positional comparison. These three replace it.
+        "loc_coverage": [],  # distinct emitted cells that are really empty / #empty
+        "loc_dup": [],       # fraction of emitted cells that repeat an earlier one
+        "loc_lcs": [],       # longest common subsequence with target order / #empty
+        "loc_wave": [],      # emitted cell sits in the earliest unfilled wave
+        "val_given_loc_acc": [],  # Correct digit AMONG steps where location matched
+        "cand_bit_acc": [],  # Per-digit accuracy of predicted candidate masks
+        "cand_set_acc": [],  # Exact candidate-SET match per empty cell (all 9 bits)
+        "cand_set_acc_changed": [],  # ...restricted to cells that changed this stage
+        "acc_complete_puzzle": []  # Accuracy of predicting correct complete puzzle
+    }
+    # Per-difficulty-level cell accuracy (levels 3..8). Diagnostic only: the
+    # curriculum no longer keys on level.
+    level_ok = {lvl: 0 for lvl in range(3, 9)}
+    level_tot = {lvl: 0 for lvl in range(3, 9)}
+
+    K = int(config.num_latent_slots)
+
+    # Per-SLOT candidate-set accuracy, i.e. per reasoning depth. Slot j holds
+    # wave snapshot j, so slot_ok[j]/slot_tot[j] is "how well is propagation
+    # block j predicted". This is the signal the depth curriculum promotes and
+    # backtracks on, replacing the old per-level accuracy.
+    slot_ok = np.zeros(max(K, 1), dtype=np.int64)
+    slot_tot = np.zeros(max(K, 1), dtype=np.int64)
+    slot_ok_ch = np.zeros(max(K, 1), dtype=np.int64)
+    slot_tot_ch = np.zeros(max(K, 1), dtype=np.int64)
+
+    # Per-stage in-set rate: was the emitted digit a *member* of that stage's
+    # candidate set? This is the promotion signal for the instance arm, where
+    # the target is one sampled assignment rather than the unique solution, so
+    # the model is right to emit any candidate. The candidate masks are read
+    # here as a metric only; nothing supervises them.
+    inset_ok = np.zeros(max(K, 1), dtype=np.int64)
+    inset_tot = np.zeros(max(K, 1), dtype=np.int64)
+
+    # Per-stage value-distribution stats, measured TEACHER-FORCED so the cell is
+    # given and the value head is scored in isolation from the model's choice of
+    # location. For a cell whose stage-s candidate set is S, the instance target
+    # is drawn uniformly from S, so the CE the trainer minimizes is
+    #     CE = -(1/|S|) sum_{d in S} log p(d)  >=  log|S|,
+    # with equality iff p == Uniform(S). Collapsing onto a single candidate
+    # sends CE to infinity. So (CE - log|S|) == KL(Uniform(S) || p) is a single
+    # number that is 0 exactly at true superposition and grows from either
+    # leakage outside S or mode collapse inside it.
+    vce_sum = np.zeros(max(K, 1), dtype=np.float64)    # sum of CE vs Uniform(S)
+    vfloor_sum = np.zeros(max(K, 1), dtype=np.float64)  # sum of log|S|
+    vmass_sum = np.zeros(max(K, 1), dtype=np.float64)   # sum of prob mass on S
+    vspread_sum = np.zeros(max(K, 1), dtype=np.float64)  # sum of H(p|S)/log|S|
+    vcnt = np.zeros(max(K, 1), dtype=np.int64)
+    vspread_cnt = np.zeros(max(K, 1), dtype=np.int64)   # only |S| >= 2 cells
+    # Same, restricted to |S| >= 2 (the cells that carry real superposition).
+    vexc_m_sum = np.zeros(max(K, 1), dtype=np.float64)
+    vmass_m_sum = np.zeros(max(K, 1), dtype=np.float64)
+    vkl_m_sum = np.zeros(max(K, 1), dtype=np.float64)
+
+    # 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))
+
+            # ---- Teacher-forced value distribution vs the candidate set ----
+            # One forward pass on the ground-truth solver-order sequence. The
+            # prefix pins down which cell each value slot refers to, so this
+            # measures the value head alone: location cannot contaminate it,
+            # and no sampling is needed because the full softmax is available.
+            if K > 0 and cand_targets is not None:
+                tf_logits, _ = run_model(input_seq, latent_vals, active_full)
+                # log_softmax over the whole vocab, matching the training CE.
+                tf_shift = tf_logits - tf_logits.max(axis=-1, keepdims=True)
+                tf_logp = tf_shift - np.log(
+                    np.exp(tf_shift).sum(axis=-1, keepdims=True))
+                for j in range(bs):
+                    si = int(start_index[j].reshape(-1)[0])
+                    n_out = 81 - si
+                    if n_out <= 0:
+                        continue
+                    base = 3 * si + K
+                    t = np.arange(n_out)
+                    v_pos = base + 3 * t + 2          # value token positions
+                    if v_pos[-1] >= config.seq_len:
+                        keep = v_pos < config.seq_len
+                        t, v_pos = t[keep], v_pos[keep]
+                        if t.size == 0:
+                            continue
+                    # Token ids are the numbers themselves: rows/cols 0..8 and
+                    # digits 1..9, so digit d lives at vocab index d.
+                    cells = (input_seq[j, base + 3 * t] * 9
+                             + input_seq[j, base + 3 * t + 1]).astype(np.int64)
+                    ok_cell = (cells >= 0) & (cells < 81)
+                    # logits at index p-1 predict the token at position p.
+                    logp_d = tf_logp[j, v_pos - 1, 1:10]        # (n, 9)
+                    for s in range(K):
+                        bits = cand_targets[j, s, np.where(ok_cell, cells, 0)]
+                        # bits <= 0 marks a clue cell, which is not supervised.
+                        sel = ok_cell & (bits > 0)
+                        if not sel.any():
+                            continue
+                        st = _cand_value_stats(logp_d[sel], bits[sel])
+                        vce_sum[s] += st["ce"]
+                        vfloor_sum[s] += st["floor"]
+                        vmass_sum[s] += st["mass"]
+                        vcnt[s] += st["count"]
+                        vspread_sum[s] += st["spread"]
+                        vspread_cnt[s] += st["spread_count"]
+                        vexc_m_sum[s] += st["excess_multi"]
+                        vmass_m_sum[s] += st["mass_multi"]
+                        vkl_m_sum[s] += st["kl_multi"]
+
+            min_start_index = int(np.min(start_index))
+            cur_input_seq = input_seq[:, :(min_start_index*3)]
+            for i in range(min_start_index * 3, config.seq_len):
+                ### In i^th iteration, i^th number in sequence will predict
+                padding = np.zeros((input_seq.shape[0],
+                                    config.seq_len - len(cur_input_seq[0])),
+                                dtype=np.int32)
+                concat_batch = np.hstack((cur_input_seq, padding))
+
+                pred_logits, _ = run_model(concat_batch, latent_vals, active_full)
+
+                # Positions < 3*start_index + K are given (clues + latent
+                # slots); the model predicts from there on. K is a multiple
+                # of 3, so the triple phase of i is unchanged.
+                if i%3 == 2:
+                    # Model predicts the value at the cell (cur_input_seq[j][i-2],
+                    # cur_input_seq[j][i-1])
+                    max_number = pred_logits[:, i-1, :].argmax(axis=-1).flatten()
+                    mask_arr = np.array(i >= (3 * start_index + K)).squeeze()
+
+                    next_number = max_number * mask_arr + (1 - mask_arr) * input_seq[:, i]
+
+                    cur_input_seq = np.hstack(
+                        (cur_input_seq, np.reshape(next_number, (-1, 1)))
+                    )
+
+                    # Iterate through all examples in batch and calculate successful
+                    # predictions of numbers
+                    for j in range(len(cur_input_seq)):
+                        if not mask_arr[j]:
+                            continue
+
+                        total_pred += 1
+                        level_tot[int(levels[j])] += 1
+                        if int(rbins[j]) in bin_tot:
+                            bin_tot[int(rbins[j])] += 1
+
+                        # Location accuracy: did the model emit the ground-truth
+                        # next cell (r,c) for this solver-order step?
+                        loc_tot += 1
+                        loc_match = (int(cur_input_seq[j][i-2]) == int(input_seq[j, i-2])
+                                     and int(cur_input_seq[j][i-1]) == int(input_seq[j, i-1]))
+                        if loc_match:
+                            loc_ok += 1
+
+                        # In-set rate per stage, scored at the ground-truth cell
+                        # so a wrong location cannot make a digit vacuously
+                        # legal. cand_targets[j, s, cell] is stage s's bitmask
+                        # under cand_slot_mode="depth" (slot s <-> stage s).
+                        if cand_targets is not None and loc_match:
+                            cell = (int(input_seq[j, i-2]) * 9
+                                    + int(input_seq[j, i-1]))
+                            v = int(cur_input_seq[j][i])
+                            for s in range(K):
+                                bits = int(cand_targets[j, s, cell])
+                                if bits <= 0:      # clue cell, not supervised
+                                    continue
+                                inset_tot[s] += 1
+                                if 1 <= v <= 9 and (bits >> (v - 1)) & 1:
+                                    inset_ok[s] += 1
+
+                        try:
+                            verify_sudoku_board(puzzle_sol[j], cur_input_seq[j][i-2], 
+                                                cur_input_seq[j][i-1], cur_input_seq[j][i])
+                        except AssertionError:
+                            # Mistake
+                            pass
+                        else:
+                            sucess_pred += 1
+                            level_ok[int(levels[j])] += 1
+                            if int(rbins[j]) in bin_ok:
+                                bin_ok[int(rbins[j])] += 1
+                            if loc_match:
+                                val_given_loc_ok += 1
+                else:
+                    # Model predicts either a row number or column number
+                    max_pos = pred_logits[:, i-1, :].argmax(axis=-1).flatten()
+                    mask = (i >= (3 * start_index + K)).squeeze()
+                    next_pos = max_pos * mask + (1 - mask) * input_seq[:, i]
+                    
+                    # pdb.set_trace()
+                    cur_input_seq = np.hstack(
+                        (cur_input_seq, np.reshape(next_pos, (-1, 1)))
+                    )
+
+            eval_metrics["acc"].append(sucess_pred * 1.0/ total_pred)
+            eval_metrics["loc_acc"].append(loc_ok * 1.0 / max(loc_tot, 1))
+            eval_metrics["val_given_loc_acc"].append(
+                val_given_loc_ok * 1.0 / max(loc_ok, 1))
+
+            def strip_latent_slots(seq, si):
+                return np.concatenate([seq[:3*si], seq[3*si + K:]])
+
+            # ---- Permutation-tolerant location diagnostics ----
+            # A cell's wave = the first stage at which its candidate set becomes
+            # a singleton, i.e. the propagation depth that determines it. Cells
+            # sharing a wave are order-interchangeable, so "did you name a cell
+            # from the earliest wave that is still unfilled" is the meaningful
+            # ordering signal; the exact index within the wave is arbitrary.
+            cov_b, dup_b, lcs_b, wav_b = [], [], [], []
+            for j in range(bs):
+                si = int(start_index[j].reshape(-1)[0])
+                pred = strip_latent_slots(cur_input_seq[j], si)
+                true = strip_latent_slots(input_seq[j], si)
+                emitted = [(int(pred[3 * k]), int(pred[3 * k + 1]))
+                           for k in range(si, 81)]
+                truth = [(int(true[3 * k]), int(true[3 * k + 1]))
+                         for k in range(si, 81)]
+                if not truth or not emitted:
+                    continue
+                n_true = len(truth)
+                cov_b.append(len(set(emitted) & set(truth)) / n_true)
+                dup_b.append(1.0 - len(set(emitted)) / len(emitted))
+                # LCS is quadratic, so sample a few examples per batch.
+                if j < 32:
+                    lcs_b.append(_lcs_len(emitted, truth) / n_true)
+
+                if cand_targets is None or K == 0:
+                    continue
+                cids = np.array([r * 9 + c for r, c in truth], dtype=np.int64)
+                bits = cand_targets[j][:, cids]                     # (K, n)
+                singleton = (bits > 0) & ((bits & (bits - 1)) == 0)
+                w = np.where(singleton.any(axis=0), singleton.argmax(axis=0), K)
+                wav_b.append(
+                    _wave_order_score(emitted, truth, w, K) / n_true)
+
+            if cov_b:
+                eval_metrics["loc_coverage"].append(float(np.mean(cov_b)))
+                eval_metrics["loc_dup"].append(float(np.mean(dup_b)))
+            if lcs_b:
+                eval_metrics["loc_lcs"].append(float(np.mean(lcs_b)))
+            if wav_b:
+                eval_metrics["loc_wave"].append(float(np.mean(wav_b)))
+
+            # ---- Print one concrete example answer the model generated ----
+            if eval_epoch == 0:
+                j = 0
+                si = int(start_index[j, 0])
+                pred = strip_latent_slots(cur_input_seq[j], si)
+                shown, n_ok, n_tot = [], 0, 0
+                for k in range(si, 81):
+                    r, c, v = int(pred[3*k]), int(pred[3*k+1]), int(pred[3*k+2])
+                    tv = int(puzzle_sol[j][r*9+c]) if (0 <= r < 9 and 0 <= c < 9) else -1
+                    ok = (0 <= r < 9 and 0 <= c < 9 and v == tv)
+                    n_tot += 1; n_ok += int(ok)
+                    if len(shown) < 12:
+                        shown.append(f"({r},{c})->{v}[true {tv}]{'ok' if ok else 'X'}")
+                if _verbose_eval():
+                    print(f"EXAMPLE (level={int(levels[j])}, k={int(k_budget[j])}): "
+                          f"model emitted {n_tot} (r,c)->v triples for the empty cells "
+                          f"(format: (row,col)->value[true T]); first 12:", flush=True)
+                    print("   ", " ".join(shown), flush=True)
+                    print(f"EXAMPLE cells-correct={n_ok}/{n_tot}  "
+                          f"valid_full_grid={valid_solution(pred)}", flush=True)
+
+                # Instance arm: emitted digit next to the deepest stage's
+                # candidate set, so it is visible whether the model is sitting
+                # inside the superposition or outside it.
+                if K > 0 and cand_targets is not None and pred_bits is None:
+                    tgt = strip_latent_slots(input_seq[j], si)
+                    shown = []
+                    for t3 in range(si, min(si + 8, 81)):
+                        r, c = int(tgt[3*t3]), int(tgt[3*t3+1])
+                        bits = int(cand_targets[j, K-1, r*9+c])
+                        cset = "".join(str(d+1) for d in range(9)
+                                       if (bits >> d) & 1)
+                        shown.append(f"(r{r},c{c})->{int(pred[3*t3+2])} "
+                                     f"in{{{cset}}}")
+                    if _verbose_eval():
+                        print(f"EXAMPLE emitted vs stage-{K} candidate set:",
+                              " ".join(shown), flush=True)
+
+                # ---- Candidate-set (multi-value) prediction for this puzzle ----
+                # Show, at the last active latent slot, predicted vs target
+                # candidate SETS for the first few empty cells. (No latent
+                # slots in the K=0 baseline, so nothing to show.)
+                if K > 0 and pred_bits is not None:
+                    kj = int(k_budget[j]) - 1
+                    def _digs(bitrow):
+                        return "".join(str(d + 1) for d in range(9) if bitrow[d])
+                    cand_shown = []
+                    for cell in range(81):
+                        if cand_targets[j, kj, cell] <= 0:   # clue / not supervised
+                            continue
+                        r, c = cell // 9, cell % 9
+                        pset = _digs(pred_bits[j, kj, cell])
+                        tset = _digs(tgt_bits[j, kj, cell])
+                        cand_shown.append(f"(r{r},c{c}) pred{{{pset}}} true{{{tset}}}")
+                        if len(cand_shown) >= 8:
+                            break
+                    if _verbose_eval():
+                        print(f"EXAMPLE candidate-set @slot{kj} (pred vs true):",
+                              " ".join(cand_shown), flush=True)
+
+            correct_eval_sudoku_puzzle = 0
+
+            for i in range(len(cur_input_seq)):
+
+                # increase correct_eval_sudoku_puzzle when the model output solution
+                # for a given puzzle is correct
+                stripped = strip_latent_slots(cur_input_seq[i], int(start_index[i, 0]))
+                correct_eval_sudoku_puzzle += valid_solution(stripped)
+
+            eval_metrics["acc_complete_puzzle"].append(
+                correct_eval_sudoku_puzzle * 1.0 / len(cur_input_seq)
+            )
+
+    per_level = {lvl: (level_ok[lvl] / level_tot[lvl] if level_tot[lvl] else -1.0)
+                 for lvl in range(3, 9)}
+    eval_metrics["per_level_acc"] = per_level
+    if _verbose_eval():
+        print("PER-LEVEL cell acc:",
+              {lvl: (f"{v:.3f}" if v >= 0 else "n/a") for lvl, v in per_level.items()},
+              flush=True)
+
+    # Per-depth candidate-set accuracy, keyed by stage (slot j -> stage j+1) so
+    # the curriculum controller can index it directly by stage number.
+    per_slot = {j + 1: (float(slot_ok[j] / slot_tot[j]) if slot_tot[j] else -1.0)
+                for j in range(K)}
+    per_slot_ch = {j + 1: (float(slot_ok_ch[j] / slot_tot_ch[j])
+                           if slot_tot_ch[j] else -1.0) for j in range(K)}
+    eval_metrics["per_slot_acc"] = per_slot
+    eval_metrics["per_slot_acc_changed"] = per_slot_ch
+
+    # Keyed by stage (slot s -> stage s+1) to match per_slot_acc.
+    per_stage_inset = {s + 1: (float(inset_ok[s] / inset_tot[s])
+                               if inset_tot[s] else -1.0) for s in range(K)}
+    eval_metrics["per_stage_inset_acc"] = per_stage_inset
+
+    # Per-stage value-distribution metrics, keyed by stage to match the above.
+    # val_excess = CE(uniform-over-candidates || model) - log|S| >= 0 is the
+    # superposition score: 0 means the model spreads exactly uniformly over the
+    # stage's candidate set, and it rises if probability leaks outside the set
+    # or collapses onto one member of it.
+    def _per_stage(num, den):
+        return {s + 1: (float(num[s] / den[s]) if den[s] else -1.0)
+                for s in range(K)}
+
+    per_stage_vce = _per_stage(vce_sum, vcnt)
+    per_stage_vfloor = _per_stage(vfloor_sum, vcnt)
+    per_stage_vexcess = {
+        s: (per_stage_vce[s] - per_stage_vfloor[s]
+            if per_stage_vce[s] >= 0 else -1.0) for s in per_stage_vce}
+    eval_metrics["per_stage_val_ce"] = per_stage_vce
+    eval_metrics["per_stage_val_floor"] = per_stage_vfloor
+    eval_metrics["per_stage_val_excess"] = per_stage_vexcess
+    eval_metrics["per_stage_val_mass"] = _per_stage(vmass_sum, vcnt)
+    eval_metrics["per_stage_val_spread"] = _per_stage(vspread_sum, vspread_cnt)
+    # Superposition metrics on |S| >= 2 cells only. excess_multi is the single
+    # number to drive to 0: it equals log(1/mass_multi) + kl_multi, so it falls
+    # only when leakage outside the set AND non-uniformity inside it both fall.
+    eval_metrics["per_stage_val_excess_multi"] = _per_stage(
+        vexc_m_sum, vspread_cnt)
+    eval_metrics["per_stage_val_mass_multi"] = _per_stage(
+        vmass_m_sum, vspread_cnt)
+    eval_metrics["per_stage_val_kl_multi"] = _per_stage(vkl_m_sum, vspread_cnt)
+    if _verbose_eval() and K > 0 and any(v >= 0 for v in per_stage_vce.values()):
+        print("PER-STAGE val_excess (0 = uniform over candidate set):",
+              {s: (f"{v:.3f}" if per_stage_vce[s] >= 0 else "n/a")
+               for s, v in per_stage_vexcess.items()}, flush=True)
+
+    if _verbose_eval() and K > 0 and any(v >= 0 for v in per_stage_inset.values()):
+        print("PER-STAGE in-set rate (emitted digit is a stage-s candidate):",
+              {s: (f"{v:.3f}" if v >= 0 else "n/a")
+               for s, v in per_stage_inset.items()}, flush=True)
+
+    per_bin = {b: (bin_ok[b] / bin_tot[b] if bin_tot[b] else -1.0)
+               for b in range(1, n_bins + 1)}
+    eval_metrics["per_bin_acc"] = per_bin
+    if _verbose_eval() and any(v >= 0 for v in per_bin.values()):
+        print("PER-ROUND-BIN cell acc:",
+              {b: (f"{v:.3f}" if v >= 0 else "n/a") for b, v in per_bin.items()},
+              flush=True)
+    if _verbose_eval() and K > 0:
+        print("PER-DEPTH cand-set acc:",
+              {s: (f"{v:.3f}" if v >= 0 else "n/a") for s, v in per_slot.items()},
+              flush=True)
+        print("PER-DEPTH cand-set acc (changed cells only):",
+              {s: (f"{v:.3f}" if v >= 0 else "n/a")
+               for s, v in per_slot_ch.items()}, flush=True)
+
+    return eval_metrics

code/wavecurriculum_run/train/main.py

@@ -0,0 +1,336 @@
+# 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.
+    # Instance arm: both gates required. mass = P(digit in raw candidate set S);
+    # spread = H(p|S)/log|S| on |S|>=2 cells. S is the wave-solver candidate
+    # set, not the filtered-instance support.
+    config.promote_acc_threshold = float(os.environ.get("SUDOKU_PROMOTE_ACC", 0.85))
+    config.promote_spread_threshold = float(
+        os.environ.get("SUDOKU_PROMOTE_SPREAD", 0.85))
+    # Location floor, applied to loc_wave. 0 = off.
+    config.promote_loc_threshold = float(
+        os.environ.get("SUDOKU_PROMOTE_LOC_WAVE", 0.0))
+    # 0 disables patience; instance arm never uses it regardless.
+    config.promote_patience_steps = int(os.environ.get("SUDOKU_PATIENCE", 0))
+    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))
+    # Superposition experiment: iterate every (puzzle, instance) pair of the
+    # pinned stage in shuffled epochs, so each instance of each puzzle is seen
+    # exactly `instance_epochs` times. 0 = old behaviour (draw a puzzle at
+    # random, then one of its assignments at random). instance_puzzles caps the
+    # puzzle count (0 = whole corpus) to shorten the experiment.
+    config.instance_epochs = int(os.environ.get("SUDOKU_INSTANCE_EPOCHS", 0))
+    config.instance_puzzles = int(os.environ.get("SUDOKU_INSTANCE_PUZZLES", 0))
+    # Preferred source: N instances per puzzle synthesized on the fly, each cell
+    # drawn uniformly from its candidate set, so the per-cell digit frequencies
+    # are equal and the CE optimum at that cell is the uniform superposition.
+    config.instance_uniform_draws = int(
+        os.environ.get("SUDOKU_UNIFORM_DRAWS", 0))
+    # 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/wavecurriculum_run/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/wavecurriculum_run/train/train_and_evaluate.py

@@ -0,0 +1,596 @@
+# 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 json
+import math
+import os
+import signal
+import socket
+import time
+import traceback
+
+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 _write_json(path, payload):
+    """Atomic JSON write so a sidecar never reads a half-written file."""
+    tmp = path + ".tmp"
+    with open(tmp, "w") as f:
+        json.dump(payload, f, indent=2, sort_keys=True)
+        f.write("\n")
+    os.replace(tmp, path)
+
+
+def _write_heartbeat(workdir, payload):
+    payload = dict(payload)
+    payload.setdefault("host", socket.gethostname())
+    payload.setdefault("job_id", os.environ.get("SLURM_JOB_ID", ""))
+    payload["unix_time"] = time.time()
+    payload["iso_time"] = time.strftime("%Y-%m-%dT%H:%M:%S%z")
+    slim = {k: payload[k] for k in
+            ("event", "step", "stage", "loss",
+             # val_acc = P(digit not in the stage-t candidate set)
+             "val_acc", "puzzle_acc",
+             # location order, tolerant to ties within a propagation wave
+             "loc_wave", "loc_coverage", "loc_lcs", "loc_dup", "loc_acc",
+             # superposition: excess 0 == uniform over the candidate set
+             "val_excess", "val_mass", "val_spread",
+             "val_excess_multi", "val_mass_multi", "val_out_multi",
+             "val_kl_multi",
+             "last_ckpt_step", "error")
+            if k in payload}
+    _write_json(os.path.join(workdir, "heartbeat.json"), slim)
+
+
+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))
+    # Applies to loc_wave (order up to wave ties), not the positional loc_acc.
+    # 0 disables the gate.
+    promote_loc_threshold = float(getattr(config, "promote_loc_threshold", 0.0))
+    # Max tolerated val_excess = CE(Uniform(S) || p) - log|S|. 0 disables.
+    promote_excess_threshold = float(
+        getattr(config, "promote_val_excess_threshold", 0.0))
+    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 = {}
+    per_stage_vexcess, per_stage_vmass, per_stage_vspread = {}, {}, {}
+    per_stage_vexc_m, per_stage_vmass_m, per_stage_vkl_m = {}, {}, {}
+    # 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
+    last_ckpt_step = start_step
+    stop_requested = {"flag": False}
+
+    def _on_stop(signum, _frame):
+        print(f"[signal] {signum} received; will save and exit after this step",
+              flush=True)
+        stop_requested["flag"] = True
+
+    signal.signal(signal.SIGTERM, _on_stop)
+    signal.signal(signal.SIGINT, _on_stop)
+    _write_heartbeat(workdir, {
+        "event": "start",
+        "step": start_step,
+        "stage": curriculum.stage,
+        "max_steps": int(config.max_steps),
+    })
+
+    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.")
+                _write_heartbeat(workdir, {
+                    "event": "nan_loss", "step": step, "stage": curriculum.stage,
+                })
+                break
+
+            if step % config.eval_every_steps == 0:
+                try:
+                    eval_metrics = evaluater.get_eval_metrics(
+                        state, eval_data_iter, p_eval_step, config)
+                except Exception:
+                    traceback.print_exc()
+                    _write_heartbeat(workdir, {
+                        "event": "eval_error",
+                        "step": step,
+                        "stage": curriculum.stage,
+                        "loss": float(metrics["loss"].mean()),
+                        "ce": float(metrics["ce_loss"].mean()),
+                        "error": traceback.format_exc()[-2000:],
+                    })
+                    print("[eval] failed; skipping this eval and continuing",
+                          flush=True)
+                    eval_metrics = None
+                if eval_metrics is not None:
+                    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_stage_vexcess = eval_metrics.pop(
+                        "per_stage_val_excess", {})
+                    per_stage_vmass = eval_metrics.pop("per_stage_val_mass", {})
+                    per_stage_vspread = eval_metrics.pop(
+                        "per_stage_val_spread", {})
+                    per_stage_vexc_m = eval_metrics.pop(
+                        "per_stage_val_excess_multi", {})
+                    per_stage_vmass_m = eval_metrics.pop(
+                        "per_stage_val_mass_multi", {})
+                    per_stage_vkl_m = eval_metrics.pop(
+                        "per_stage_val_kl_multi", {})
+                    eval_metrics.pop("per_stage_val_ce", {})
+                    eval_metrics.pop("per_stage_val_floor", {})
+                    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
+                    def _s(d):
+                        return round(float(d.get(curriculum.stage, -1.0)), 4)
+                    # val_acc = P(digit not in stage-t candidate set),
+                    # teacher-forced so location cannot contaminate it.
+                    # 0 = all mass inside S; chance at stage 0 is ~0.59.
+                    mass = _s(per_stage_vmass)
+                    val_out = round(1.0 - mass, 4) if mass >= 0 else -1.0
+                    mass_m = _s(per_stage_vmass_m)
+                    val_out_m = round(1.0 - mass_m, 4) if mass_m >= 0 else -1.0
+                    print(step,
+                          "stage", curriculum.stage,
+                          "loss", round(float(metrics["loss"].mean()), 4),
+                          "val_acc", val_out,
+                          "loc_wave", _m("loc_wave"),
+                          "loc_cov", _m("loc_coverage"),
+                          # superposition on |S|>=2 cells: excess = leak + kl
+                          "excess*", _s(per_stage_vexc_m),
+                          "out*", val_out_m,
+                          "kl*", _s(per_stage_vkl_m),
+                          "val_spread", _s(per_stage_vspread),
+                          flush=True)
+                    _write_heartbeat(workdir, {
+                        "event": "eval",
+                        "step": int(step),
+                        "stage": int(curriculum.stage),
+                        "loss": round(float(metrics["loss"].mean()), 6),
+                        "ce": round(float(metrics["ce_loss"].mean()), 6),
+                        "val_acc": val_out,
+                        "puzzle_acc": _m("acc_complete_puzzle"),
+                        # Permutation-tolerant location signals; loc_acc is the
+                        # old positional match, kept only for continuity.
+                        "loc_wave": _m("loc_wave"),
+                        "loc_coverage": _m("loc_coverage"),
+                        "loc_lcs": _m("loc_lcs"),
+                        "loc_dup": _m("loc_dup"),
+                        "loc_acc": _m("loc_acc"),
+                        # Superposition: 0 excess == uniform over the stage's
+                        # candidate set; spread 1.0 == no mode collapse.
+                        "val_excess": _s(per_stage_vexcess),
+                        "val_mass": _s(per_stage_vmass),
+                        "val_spread": _s(per_stage_vspread),
+                        # |S|>=2 cells only: excess_multi == log(1/mass_multi)
+                        # + kl_multi, so it is the single number that falls only
+                        # when leakage and non-uniformity both fall.
+                        "val_excess_multi": _s(per_stage_vexc_m),
+                        "val_mass_multi": mass_m,
+                        "val_out_multi": val_out_m,
+                        "val_kl_multi": _s(per_stage_vkl_m),
+                        "last_ckpt_step": int(last_ckpt_step),
+                    })
+                    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_wave")
+                        if instance_mode:
+                            # Frontier = probability mass the value head puts on
+                            # stage t's candidate set, measured teacher-forced so
+                            # it does not depend on the model's choice of order.
+                            # This is the soft form of the old in-set rate, which
+                            # was conditioned on an exact positional location
+                            # match and so was estimated from ~4% of steps.
+                            # Frontier = exp(-val_excess_multi) on stage t, i.e.
+                            # the superposition score mapped into (0, 1] so the
+                            # existing ">= threshold" machinery still applies.
+                            # val_excess_multi = log(1/mass) + KL(Uniform(S)||q)
+                            # on |S| >= 2 cells, so it reaches 0 (score 1.0)
+                            # only when the model puts all its mass on the
+                            # candidate set AND spreads uniformly over it.
+                            exc_now = per_stage_vexc_m.get(t, -1.0)
+                            frontier_acc = (float(np.exp(-exc_now))
+                                            if exc_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
+                        # Location gate. The upstream code never compared the
+                        # emitted (r,c) to the target order at all, and the
+                        # earlier gate on the positional loc_acc was
+                        # unsatisfiable (0.70 required, ~0.04 observed), so every
+                        # promotion fired on the patience cap alone. It now keys
+                        # on loc_wave, which credits any cell from the earliest
+                        # unfilled propagation wave, and defaults to off until
+                        # the reachable range is known from a run.
+                        loc_ready = (promote_loc_threshold <= 0.0
+                                     or loc_now >= promote_loc_threshold)
+                        excess_now = per_stage_vexcess.get(t, -1.0)
+                        excess_ready = (promote_excess_threshold <= 0.0
+                                        or (0.0 <= excess_now
+                                            <= promote_excess_threshold))
+                        gates_ok = loc_ready and excess_ready
+                        hit_threshold = (measured and gates_ok
+                                         and frontier_acc >= promote_threshold)
+                        stalled = (measured and gates_ok 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}_")
+                                last_ckpt_step = step
+                                _write_json(os.path.join(workdir, "ckpt_ready.json"), {
+                                    "event": "stage",
+                                    "step": int(step),
+                                    "stage": int(curriculum.stage),
+                                    "reason": reason,
+                                    "workdir": workdir,
+                                    "stage_ckpt_dir": stage_ckpt_dir,
+                                    "keep_snapshot": True,
+                                })
+                                _write_heartbeat(workdir, {
+                                    "event": "promote",
+                                    "step": int(step),
+                                    "stage": int(curriculum.stage),
+                                    "reason": reason,
+                                    "graduated_acc": round(float(frontier_acc), 6),
+                                    "last_ckpt_step": int(last_ckpt_step),
+                                })
+
+            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
+                )
+                last_ckpt_step = step
+                _write_json(os.path.join(workdir, "ckpt_ready.json"), {
+                    "event": "periodic",
+                    "step": int(step),
+                    "stage": int(curriculum.stage),
+                    "workdir": workdir,
+                    "keep_snapshot": (step % 50000 == 0),
+                })
+                _write_heartbeat(workdir, {
+                    "event": "ckpt",
+                    "step": int(step),
+                    "stage": int(curriculum.stage),
+                    "last_ckpt_step": int(last_ckpt_step),
+                })
+
+            if stop_requested["flag"]:
+                print(f"[signal] stopping at step {step}", flush=True)
+                if config.save_checkpoint:
+                    checkpoints.save_checkpoint_multiprocess(
+                        workdir, jax_utils.unreplicate(state), step,
+                        keep=ckpt_keep, overwrite=True)
+                    last_ckpt_step = step
+                    _write_json(os.path.join(workdir, "ckpt_ready.json"), {
+                        "event": "signal",
+                        "step": int(step),
+                        "stage": int(curriculum.stage),
+                        "workdir": workdir,
+                        "keep_snapshot": True,
+                    })
+                _write_heartbeat(workdir, {
+                    "event": "stopped",
+                    "step": int(step),
+                    "stage": int(curriculum.stage),
+                    "last_ckpt_step": int(last_ckpt_step),
+                })
+                break
+
+        # Final checkpoint at the end of training
+        if config.save_checkpoint and not stop_requested["flag"]:
+            checkpoints.save_checkpoint_multiprocess(
+                workdir, jax_utils.unreplicate(state), config.max_steps,
+                keep=ckpt_keep, overwrite=True)
+            _write_json(os.path.join(workdir, "ckpt_ready.json"), {
+                "event": "final",
+                "step": int(config.max_steps),
+                "stage": int(curriculum.stage),
+                "workdir": workdir,
+                "keep_snapshot": True,
+            })
+        _write_heartbeat(workdir, {
+            "event": "done",
+            "step": int(last_ckpt_step),
+            "stage": int(curriculum.stage),
+            "last_ckpt_step": int(last_ckpt_step),
+        })
+
+  

code/wavecurriculum_run/train/train_backtrack.py

@@ -0,0 +1,763 @@
+"""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_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", {})
+                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
+
+                # 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
+                        hit_threshold = measured and frontier_acc >= promote_threshold
+                        stalled = (measured 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/wavecurriculum_run/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/wavecurriculum_run/verify_uniform_loader.py

@@ -0,0 +1,132 @@
+"""Verify the on-the-fly uniform instance sampler on the real stage-1 masks.
+
+Replicates SudokuDataset.uniform_instance_values / apply_instance with numpy
+only, then checks the three things the experiment depends on:
+  1. the clue block is untouched, so the input prompt is fixed per puzzle
+  2. every drawn digit lies in that cell's candidate set
+  3. over N draws each candidate of each cell turns up a near-equal number of
+     times, and at least 5 times
+"""
+import numpy as np
+
+PUZ = ("/scratch/users/gatmiry/llm-reasoning-logic-puzzles/sudoku-code/"
+       "datasets/train_sudoku_puzzles.npy")
+CAND = "/tmp/sudoku_s12/train_cand_masks.npy"
+STAGE = 0
+N_DRAWS = 64
+K, LATENT_ID = 12, 10
+
+
+def load_inputs(n):
+    raw = np.load(PUZ, mmap_mode="r")
+    rows = np.asarray(raw[:n]).astype(np.int64)
+    return np.delete(rows[:, 1:], np.arange(81) * 4 + 3, axis=1), rows[:, 0]
+
+
+def uniform_instance_values(mask, rng):
+    """Same argmax-of-random-keys draw as the loader."""
+    m = mask.astype(np.int64)
+    bits = ((m[:, None] >> np.arange(9)) & 1).astype(np.float64)
+    keys = rng.random_sample((81, 9)) * bits
+    vals = (keys.argmax(1) + 1).astype(np.int8)
+    return np.where(m > 0, vals, 0)
+
+
+def apply_instance(seq, vals):
+    seq = seq.copy()
+    cells = seq[0::3] * 9 + seq[1::3]
+    new = vals[cells]
+    seq[2::3] = np.where(new > 0, new, seq[2::3])
+    return seq
+
+
+def digits_of(m):
+    return [d for d in range(1, 10) if int(m) & (1 << (d - 1))]
+
+
+def main():
+    n_show = 500
+    inputs, si_all = load_inputs(n_show)
+    masks = np.load(CAND, mmap_mode="r")
+    rng = np.random.RandomState(0)
+
+    p = 3
+    si, seq = int(si_all[p]), inputs[p]
+    mask = np.asarray(masks[p, STAGE])
+    tr = seq.reshape(-1, 3)
+    empty = (tr[si:, 0] * 9 + tr[si:, 1])
+
+    drawn = [apply_instance(seq, uniform_instance_values(mask, rng))
+             for _ in range(N_DRAWS)]
+
+    print(f"puzzle p={p}: {si} clues, {len(empty)} empty cells, "
+          f"{N_DRAWS} uniform draws")
+
+    clue_ok = all(np.array_equal(d[:3 * si], seq[:3 * si]) for d in drawn)
+    print(f"  clue block identical in all {N_DRAWS} draws: {clue_ok}")
+
+    bad = sum(int(d[2::3][j + si]) not in digits_of(mask[empty[j]])
+              for d in drawn for j in range(len(empty)))
+    print(f"  drawn digits outside the candidate set: {bad}")
+
+    blk = np.stack([d[2::3] for d in drawn])[:, si:]     # (N, n_empty)
+
+    print(f"\ninstance = one digit for every empty cell. First 10 cells:")
+    print("  cells:  " + "  ".join(f"({c // 9},{c % 9})" for c in empty[:10]))
+    for i in range(4):
+        print(f"  i{i + 1}:     " + "     ".join(
+            f"{int(x)}" for x in blk[i, :10]))
+    print("  S:      " + "  ".join(f"{digits_of(mask[c])}" for c in empty[:4])
+          + " ...")
+
+    print(f"\nreading DOWN a cell's column over {N_DRAWS} draws:")
+    print(f"  {'cell':>8}  {'|S|':>3}  {'S':<16}{'counts':<34}min")
+    shown = 0
+    worst = 10 ** 9
+    for j, c in enumerate(empty):
+        S = digits_of(mask[c])
+        if len(S) < 2:
+            continue
+        v, k = np.unique(blk[:, j], return_counts=True)
+        cnt = {int(a): int(b) for a, b in zip(v, k)}
+        worst = min(worst, min(cnt.get(d, 0) for d in S))
+        shown += 1
+        if shown > 8:
+            continue
+        cs = " ".join(f"{d}:{cnt.get(d, 0)}" for d in S)
+        print(f"  ({c // 9},{c % 9})     {len(S):>3}  {str(S):<16}{cs:<34}"
+              f"{min(cnt.get(d, 0) for d in S)}")
+    print(f"  ... {shown} multi-candidate cells; rarest candidate anywhere in "
+          f"this puzzle appeared {worst}x (need >=5)")
+
+    print(f"\nover {n_show} puzzles:")
+    ge5 = tot = 0
+    spreads = []
+    for q in range(n_show):
+        m = np.asarray(masks[q, STAGE])
+        b = np.stack([uniform_instance_values(m, rng) for _ in range(N_DRAWS)])
+        for c in range(81):
+            S = digits_of(m[c])
+            if len(S) < 2:
+                continue
+            v, k = np.unique(b[:, c], return_counts=True)
+            cnt = {int(a): int(b_) for a, b_ in zip(v, k)}
+            tot += len(S)
+            ge5 += sum(1 for d in S if cnt.get(d, 0) >= 5)
+            pr = np.array([cnt.get(d, 0) for d in S], dtype=float)
+            pr /= pr.sum()
+            nz = pr[pr > 0]
+            spreads.append(float(-(nz * np.log(nz)).sum()) / np.log(len(S)))
+    print(f"  (cell, candidate) pairs seen >=5 times: {ge5 / tot:.4f}")
+    print(f"  mean spread H(p)/log|S|: {np.mean(spreads):.4f}  "
+          f"(1.0 = uniform superposition)")
+
+    full = np.concatenate([drawn[0][:3 * si],
+                           np.full(K, LATENT_ID, dtype=drawn[0].dtype),
+                           drawn[0][3 * si:]])
+    print(f"\ntoken sequence: {len(full)} = {3 * si} clue + {K} latent + "
+          f"{3 * (81 - si)} output")
+
+
+if __name__ == "__main__":
+    main()