# Run the noisy-latent Sudoku trainer

Data lives on Hugging Face: [Avra98/Sudoku_superposition](https://huggingface.co/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

```bash
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

```bash
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
bash code/run_experiment.sh
```

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

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

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

```text
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`):

```bash
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/`.