# 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). The round-bin is 0 unless the round-count data # curriculum is enabled; the counts are zero on the train split, which # never reads them. K = int(getattr(config, "num_latent_slots", 0)) output_types = (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]), ) # 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. """ 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 _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 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, ) 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)