Run the noisy-latent Sudoku trainer

Data lives on Hugging Face: Avra98/Sudoku_superposition. This folder is the trainer (train/). Do not regenerate masks or instances unless you changed a generation flag.

1. Download the 11 required files

pip install -U "huggingface_hub[hf_xet]"

huggingface-cli download Avra98/Sudoku_superposition \
  --include "data/train_sudoku_puzzles.npy" \
  --include "data/test_sudoku_puzzles.npy" \
  --include "data/train_cand_masks.npy" \
  --include "data/test_cand_masks.npy" \
  --include "data/train_assignments.npy" \
  --include "data/train_starts.npy" \
  --include "data/train_counts.npy" \
  --include "data/test_assignments.npy" \
  --include "data/test_starts.npy" \
  --include "data/test_counts.npy" \
  --include "data/test_bigq300.npz" \
  --include "code/*" \
  --local-dir ./Sudoku_superposition

Skip *_index.npy, *_meta.npy, *_solution.npy.

2. Install trainer deps

pip install jax flax optax ml_collections clu absl-py tensorflow wandb

GPU JAX must match your CUDA. CPU will start but is too slow for the full run.

3. Run

From Sudoku_superposition/ (HF layout: data/ + code/):

bash code/run_experiment.sh

From this GitHub tree (sudoku-wavecurriculum/), after the HF data/ folder exists:

SUDOKU_DATA_DIR=/path/to/Sudoku_superposition/data bash run_experiment.sh

The script checks that every required file exists, then runs:

python -m train.main --workdir=../runs/noise --exp_name=noise

If you rebuilt instances

Keep the premade masks in data/. Point only the instance dir at the new folder (must contain {train,test}_{assignments,starts,counts}.npy):

SUDOKU_INSTANCE_DIR=/path/to/datasets_superposition bash code/run_experiment.sh

If that folder sits at Sudoku_superposition/datasets_superposition/ and is complete, the script picks it automatically.

Optional env knobs

var	default	meaning
`SUDOKU_LATENT_NOISE`	`1`	noisy latents on
`SUDOKU_MC_SAMPLES`	`8`	Monte-Carlo draws
`SUDOKU_LATENT_SLOTS`	`12`	must match the 12-stage masks
`SUDOKU_AUX_WEIGHT`	`0`	candidate-head off (instance CE)
`SUDOKU_START_STAGE`	`1`	set `4` to skip early curriculum
`SUDOKU_MINIBATCH`	`64`	global batch; must divide GPU count
`WORKDIR`	`runs/noise`	checkpoints + logs

Do not run staged_candidate_gen.py or build_superposition_dataset.py to start training. Those files are already in data/.