Avra98/Sudoku_superposition

Sudoku superposition instances

Concrete board assignments for a 12-stage Sudoku latent curriculum. Each stage keeps a candidate set per cell; this dataset materializes those sets as ordinary (row, col, value) sequences so training is standard next-token cross-entropy (no multi-hot candidate head).

Repo: Avra98/Sudoku_superposition

Dataset

split	puzzles	instances	mean / puzzle
train	1,804,462	89,236,838	49.45
test	99,999	4,945,532	49.46

Mean instances per stage (train, stage 0 → 11): 6.28, 5.96, 5.65, 5.34, 5.02, 4.68, 4.32, 3.88, 3.33, 2.58, 1.41, 1.00

Stage 11 is the unique solution. No empty (puzzle, stage) pair.

Files (data/)

file	shape	dtype	role
{split}_assignments.npy	(M, 81)	uint8	one full board per instance, cell r*9+c
{split}_starts.npy	(N, 12)	int32	first row in assignments for (puzzle, stage)
{split}_counts.npy	(N, 12)	uint8	number of instances for (puzzle, stage)
{split}_index.npy	(M, 3)	int32	[puzzle_idx, stage, k] (optional; starts/counts are enough)

The trainer only needs assignments, starts, and counts.

Puzzle clue/solution arrays are not in this repo (they are the original Sudoku npy files). Candidate masks used to build the instances live in datasets_multicandidate_s12/.

How a training example is built

1. Take the usual solver-order sequence: clue triples, then K latent placeholders, then empty-cell triples.
2. Sample one instance for the current curriculum stage.
3. Rewrite values only. Location order stays solver-order.
4. Predict the output triples with softmax CE. Latent slots are not predicted.

Curriculum stage t (1..12) trains on stage-(t-1) instances. Stage 12 targets the unique solution.

Code (code/)

• code/train/ — JAX trainer (data.py, trainer.py, evaluater.py, train_and_evaluate.py, train_backtrack.py, main.py, model.py)
• code/build_superposition_dataset.py — instance generator
• code/build_instance_offsets.py — starts / counts tables
• code/superposition_instances.py — per-puzzle instance sampler
• code/sbatch_instance_latent.sh — Slurm launch (feanor / H200)

Set SUDOKU_INSTANCE_DIR to the data/ folder (or a local copy).

from huggingface_hub import snapshot_download
snapshot_download("Avra98/Sudoku_superposition", local_dir="Sudoku_superposition")