code/train/data.py

# 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, per-slot
    # instance-digit counts, instance-valued sequences). The round-bin is 0
    # unless the round-count data curriculum is enabled; the counts and the
    # instance sequences are zero on the train split, which never reads them.
    K = int(getattr(config, "num_latent_slots", 0))
    R = int(getattr(config, "eval_instance_reads", 4))
    output_types = (tf.int32, tf.int32, 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]),
        tf.TensorShape([K, config.block_size, 9]),
        tf.TensorShape([R, config.seq_len]),
    )

    # 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)
    tf_ds = tf_ds.prefetch(tf.data.AUTOTUNE)
    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._load_gate_q()
        self._build_instance_epoch_list()

    def _load_gate_q(self):
        """Optional large-sample instance pool for the promotion gate only.

        The stored pool is coverage-minimal: `instances_for_stage` runs
        `while uncovered and ...` with require_new rejecting redundant draws,
        so it stops as soon as every (cell, candidate) pair has appeared once.
        That gives ~5 grids per puzzle (1.4 by stage 11) and it deliberately
        picks rare candidates, which flattens the per-cell histogram toward
        uniform. The resulting q is not a frequency estimate: at stage 4 it
        reports q(true) = 0.35 against 1/|S| = 0.35, while re-sampling the same
        generator 200 times per cell gives q(true) = 0.59. The gate spends
        log|S| - H(q) as its headroom, so this is not a detail -- it is the
        difference between 0.15 nats of headroom and 0.57.

        Set SUDOKU_GATE_Q to an npz from gen_bigq.py to score the gate against
        the large-sample q instead. It covers a fixed prefix of the split, so
        eval is restricted to those puzzles; a few hundred is ample for a
        statistic averaged over ~40 candidate cells each.
        """
        self.gate_q = None
        p = getattr(self.config, "gate_q_path", None)
        if not p or self.train:
            return
        z = np.load(p)
        self.gate_q = z["counts"]                      # (P, S, 81, 9)
        self.gate_grids = z["grids"]                   # (P, S, R, 81)
        P = int(self.gate_q.shape[0])
        self.eval_inputs = self.eval_inputs[:P]
        self.eval_puzzles = self.eval_puzzles[:P]
        self.eval_start_index = self.eval_start_index[:P]
        self.eval_levels = self.eval_levels[:P]
        self.level_index = self._build_level_index(self.eval_levels)
        print(f"[gate-q] {p}: {P} puzzles, "
              f"{int(z['accepted'].mean())} draws/cell mean; eval restricted "
              f"to those puzzles", flush=True)

    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.
        """
        ov = getattr(self, "instance_stage_override", None)
        if ov is not None:
            S = int(getattr(self, "inst_stages", 0)) or int(self.num_stages)
            return int(np.clip(int(ov), 0, max(S - 1, 0)))
        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 sample_instance_batch(self, bs, inst_stage):
        """Numpy batch whose value tokens are instances of 0-based `inst_stage`.

        Used by depth-mix so each depth can pull its own plane without going
        through tf.data. The TF pipeline shuffles and prefetches across a
        shared generator: flipping curriculum.stage under it mixed planes and
        rebuilt the 6-11M pair list on every depth of every step.
        """
        n = len(self.train_puzzles) if self.train else len(self.eval_puzzles)
        n_puz = int(getattr(self.config, "instance_puzzles", 0))
        n = n if n_puz <= 0 else min(n_puz, n)
        prev = getattr(self, "instance_stage_override", None)
        self.instance_stage_override = int(inst_stage)
        try:
            idxs = self.rng.randint(n, size=int(bs))
            seqs, puzzles, starts, levels, cands = [], [], [], [], []
            for idx in idxs:
                item = self.__getitem__(int(idx))
                seq, puzzle, start_index, lvl, cand = item[:5]
                seqs.append(seq)
                puzzles.append(puzzle)
                starts.append(start_index)
                levels.append(lvl)
                cands.append(cand)
        finally:
            if prev is None:
                self.instance_stage_override = None
            else:
                self.instance_stage_override = prev
        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 _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_stages(self, level):
        """Stage index backing each of the K latent slots, for one example.

        Factored out of _slot_stage_targets so the candidate bitmasks and the
        instance-frequency counts below are guaranteed to describe the same
        stage at the same slot.
        """
        K = self.num_latent_slots
        S = self.num_stages
        k = self.slot_budget(level)
        depth_mode = getattr(self.config, "cand_slot_mode", "level") == "depth"
        out = []
        for j in range(K):
            if j >= k:
                # Inactive slot: masked out of the loss, value is irrelevant.
                out.append(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.
                out.append(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).
                out.append(int(round(j * (S - 1) / max(k - 1, 1))))
        return out

    def _slot_stage_qcounts(self, idx, level, clue_cells=None):
        """Return (K, 81, 9) int32 instance-digit counts, one plane per slot.

        Entry [j, c, d-1] is how many of this puzzle's stored stage-s instances
        put digit d in cell c, where s is the stage behind slot j. Normalizing
        over d gives q, the post-constraint distribution the training data
        actually teaches at that cell -- as opposed to Uniform(S) over the raw
        pre-constraint candidate set, which the constraints have already pruned.

        Eval needs this because CE(q || p) is the only value statistic whose
        floor, H(q), is strictly below log|S| when the constraints bite. CE
        against Uniform(S) is >= log|S| for every p by Gibbs, so comparing that
        to log|S| tests nothing. Built for eval only; train never reads it.
        """
        K = self.num_latent_slots
        out = np.zeros((K, 81, 9), dtype=np.int32)
        if self.instances is None or K == 0 or self.train:
            return out
        gq = getattr(self, "gate_q", None)
        if gq is not None and idx < gq.shape[0]:
            for j, s in enumerate(self._slot_stages(level)):
                out[j] = gq[idx, min(s, gq.shape[1] - 1)]
            if clue_cells is not None and len(clue_cells) > 0:
                out[:, clue_cells, :] = 0
            return out
        for j, s in enumerate(self._slot_stages(level)):
            s = min(s, self.inst_stages - 1)
            n = int(self.inst_counts[idx, s])
            if n <= 0:
                continue
            lo = int(self.inst_starts[idx, s])
            rows = np.asarray(self.instances[lo:lo + n], dtype=np.int64)  # (n,81)
            # Digits are 1..9; 0 marks "not assigned by this instance".
            valid = (rows >= 1) & (rows <= 9)
            cells = np.broadcast_to(np.arange(81)[None, :], rows.shape)
            np.add.at(out[j], (cells[valid], rows[valid] - 1), 1)
        if clue_cells is not None and len(clue_cells) > 0:
            out[:, clue_cells, :] = 0
        return out

    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)
        stages = self.cand_masks[idx].astype(np.int32)   # (S, 81)
        out = np.zeros((K, 81), dtype=np.int32)
        for j, s in enumerate(self._slot_stages(level)):
            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)
        q_counts = self._slot_stage_qcounts(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),
            q_counts,
            self._instance_seqs(idx, si),
        )

    def _instance_seqs(self, idx, si):
        """R copies of the sequence, values taken from R different instances.

        The solution-valued sequence above is right for accuracy and for free
        generation, but wrong for scoring CE(q||p): reading p at v_pos - 1 on
        it conditions the prediction on the TRUE digit of every earlier empty
        cell, so the model resolves the current cell by ordinary elimination
        and is then charged for not matching q, which conditions on nothing of
        the sort. Training never does this -- it always rewrites the values
        with an instance -- so the gate was also measuring a distribution the
        model was never trained on. These sequences restore the match; the
        caller averages p over the R of them to recover a marginal comparable
        to q.
        """
        R = int(getattr(self.config, "eval_instance_reads", 4))
        out = np.zeros((R, self.config.seq_len), dtype=np.int32)
        if self.train or self.instances is None or R <= 0:
            return out
        stage = self.instance_stage()
        base = self.eval_inputs[idx, :]
        gq = getattr(self, "gate_q", None)
        if gq is not None and idx < gq.shape[0]:
            grids = self.gate_grids[idx, min(stage, gq.shape[1] - 1)]
            for r in range(R):
                out[r] = np.asarray(self.insert_latent_slots(
                    self._apply_values(base, grids[r % len(grids)]), si),
                    dtype=np.int32)
            return out
        stage = min(stage, self.inst_stages - 1)
        n = int(self.inst_counts[idx, stage])
        if n <= 0:
            return out
        lo = int(self.inst_starts[idx, stage])
        for r in range(R):
            # Cycle when the pool is smaller than R, which it is at the deepest
            # stages (mean 1.4 instances per puzzle by stage 11).
            row = lo + (r % n)
            seq = self.insert_latent_slots(
                self.apply_instance(base, idx, stage, inst_row=row), si)
            out[r] = np.asarray(seq, dtype=np.int32)
        return out

    def _apply_values(self, seq, vals):
        """apply_instance for an assignment that is passed in rather than
        looked up in the stored table."""
        seq = np.asarray(seq).copy()
        cells = seq[0::3].astype(np.int64) * 9 + seq[1::3].astype(np.int64)
        new = np.asarray(vals)[cells].astype(seq.dtype)
        seq[2::3] = np.where(new > 0, new, seq[2::3])
        return seq

    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)