giant/issues.md

# GIANT — architecture issues

Software-engineering review of the `giant` codebase, conducted 2026-08-12 on branch
`v0.3.0-stage2-autoregressive` at commit `55332db`.

This document is written to be read **without the context of the review conversation**.
Each issue states what the code does today, why that is a problem, how to verify the
claim independently, what the fix looks like, and — importantly — what the fix should
*not* touch. Line numbers are accurate as of `55332db`; if they have drifted, the
accompanying code excerpts and `grep` commands will still locate the code.

## Baseline: what is already good

Read this first, so the issues below are calibrated correctly. **This is a healthy
codebase.** The problems listed are about structural resilience as the v0.3.0
architecture matrix grows, not about rot or breakage.

- `pytest` — **725 tests pass in ~18 s**. Fast enough that there is no excuse for not
  running it on every change.
- `uv run ruff check .`, `uv run ruff format --check .`, `uv run ty check .` — all clean,
  and all three are enforced in CI (`.gitea/workflows/ci.yml`).
- Zero `TODO` / `FIXME` / `XXX` / `HACK` markers in `giant/` or `scripts/`.
- The internal import graph is **acyclic** with clean layering:
  `constants → data → model → training → pipeline → cli`, and `analysis` almost fully
  independent of the rest.
- Line coverage is **96–99.6 % on every core module** (`network.py` 98.6 %,
  `trainers.py` 99.5 %, `rollout.py` 99.6 %, `transforms.py` 96.2 %, `pipeline.py`
  97.1 %, `analysis/catalog.py` 99.6 %). The single exception is `cli.py` at 35.8 % —
  see Issue 4.
- Comment and docstring quality is unusually high. Docstrings routinely explain *why*
  a decision was made, not just what the code does, and several call out their own
  known limitations honestly. **Preserve this when refactoring.** A refactor that
  deletes the reasoning in `giant/config.py`'s `DEFAULT_CONFIG` comments or
  `giant/analysis/router_gating.py`'s module docstring is a net loss even if the code
  gets shorter.
- `giant/analysis/` is the best-designed subsystem in the repo and should be treated as
  the template the rest of the codebase moves toward. See Issue 11.

## Issue index

| # | Issue | Severity | Effort | Status |
|---|---|---|---|---|
| 1 | Config defaults are declared twice; `DEFAULT_CONFIG` and consumers already disagree | **High** | Medium | **Fixed** (`9bf5874`) |
| 2 | No unknown-key validation — a typo in `config.toml` silently trains the wrong model | **High** | Small | **Fixed** |
| 3 | `cli.py:train()` is a 58-parameter, 510-line fat controller | **High** | Medium | **Fixed** (`2bfb1ab`) |
| 4 | `cli.py` is at 35.8 % coverage and holds untested override-precedence logic | **High** | Medium | **Fixed** (`2bfb1ab`, partial — see status note) |
| 5 | Inference bootstrap is duplicated verbatim between `predict` and `rollout` | **High** | Small | **Fixed** |
| 6 | Two independent v0.2→v0.3 migration surfaces encode the same knowledge | Medium | Medium | Open |
| 7 | Positional tuple contracts (9-tuple, 7-tuple) between data, model and training layers | Medium | Small | Open |
| 8 | `network.py` is 1745 lines holding three distinct modules | Medium | Small | Open |
| 9 | `scripts` is published as a top-level distribution package | Medium | Small | Open |
| 10 | `torch.load(weights_only=False)` — checkpoints are arbitrary pickles | Low | Medium | Open |
| 11 | Minor: `echo=print` threading, `particles → data.loader` layering | Low | Small | Open |

---

## Issue 1 — Config defaults are declared twice, and the two declarations already disagree

> **Status: Fixed, commit `9bf5874` on `v0.3.0-stage2-autoregressive`.** `giant/config.py` now
> declares a hierarchy of frozen dataclasses (`GiantConfig` and its nested blocks, following
> `StageSpec`'s existing style) as the single source of truth; `DEFAULT_CONFIG` is generated from
> `GiantConfig().to_dict()` rather than hand-maintained, and `build_models`/`build_critics`
> (`network.py`) and `StageSpec.from_config` (`trainers.py`) convert their dict input into these
> dataclasses at the top of each function and read attributes instead of duplicating
> `.get(key, literal)` defaults. Both drifted keys (`decoder`, `particle_type.target`) now resolve
> to `DEFAULT_CONFIG`'s documented v0.3.0 values (`"autoregressive"`/`"onehot"`) unconditionally.
> Router/`n_sec` sub-blocks keep a dict-shaped `extra` catch-all for their genuinely dynamic keys
> (composed-router `axis{i}_*`, runtime-seeded `centers_init`, legacy `legacy_owner`) rather than
> being fully typed — see the "Recommended fix" section below, which this diverges from slightly
> on that one point. Fixing the fallback surfaced two existing callers that had been silently
> depending on the old (wrong) default — a `tests/test_train.py` fixture and
> `scripts/warm_setup_cache.py`'s hand-built minimal config (now merged against `DEFAULT_CONFIG`
> instead of hand-rolled) — both fixed in the same commit. Regression tests added in
> `tests/test_config.py`, `tests/test_network.py`, `tests/test_train.py`, including one pinning
> `GiantConfig().to_dict() == DEFAULT_CONFIG` so this class of drift can't recur silently.
> Everything below this point describes the pre-fix state and is kept for historical context.

**Severity: High. Effort: Medium. Risk if unfixed: silent wrong-model training.**

**Location:** `giant/config.py:32` (`DEFAULT_CONFIG`) versus `giant/model/network.py:1570-1693`
(`build_models`), `giant/training/trainers.py:130-175` (`StageSpec.from_config`),
`giant/pipeline.py:381`.

### What the code does today

`DEFAULT_CONFIG` in `giant/config.py:32` is a fully-specified nested dict. Every key has
a value and most have an explanatory comment. It is, on paper, the single source of
truth for configuration defaults.

The consumers do not treat it that way. They re-declare the same defaults inline via
`dict.get(key, default)`:

```python
# giant/model/network.py, inside build_models
cond_out_dim = conditioning.get("out_dim", 128)
...
hidden_dim=s1cfg.get("hidden_dim", 256),
n_res_blocks=s1cfg.get("n_res_blocks", 6),
dropout=s1cfg.get("dropout", 0.0),
time_dim=gen_sub.get("time_dim", 64),
```

```python
# giant/training/trainers.py, inside StageSpec.from_config
lambda_weight=stage_cfg.get("lambda", 1.0),
n_sec_lambda=cfg["stage2_model"].get("n_sec", {}).get("lambda", 0.1),
ddpm_n_steps=stage_cfg.get("ddpm", {}).get("n_steps", 1000),
n_critic=wgan_cfg.get("n_critic", 5),
gp_weight=wgan_cfg.get("gp_weight", 10.0),
```

There are **132 `.get("key", default)` call sites across `giant/`**, concentrated in
`network.py` (55) and `trainers.py` (25). Each one is an independent copy of a value
that `DEFAULT_CONFIG` also declares.

### Why it's a problem

Because `cfg` is typed `dict` (i.e. `dict[str, Any]`), neither `ty` nor the tests can
see that the two declarations are supposed to agree. When they drift, nothing fails —
the model is simply built differently than the config file says.

**This has already happened.** Two defaults currently disagree:

| Key | `DEFAULT_CONFIG` says | `build_models` fallback says |
|---|---|---|
| `stage2_model.decoder` | `"autoregressive"` | `"one_shot"` (`network.py:1632`) |
| `stage2_model.particle_type.target` | `"onehot"` | `"physical"` (`network.py:1644`, `trainers.py:143`, `pipeline.py:381`) |

Both are *architecturally load-bearing*: `decoder` selects between `Stage2Autoregressive`
and `Stage2OneShot` — two different networks — and `particle_type.target` selects between
categorical class logits and continuous `(log-mass, charge)` regression, which is the
exact axis the v0.3.0 redesign was created to change (see `CLAUDE.md`, "v0.3.0 —
Stage-2 autoregressive redesign").

### Is it currently a live bug?

**No — it is currently latent.** Two things save it today:

1. The training path always passes a fully-merged config. `giant/pipeline.py:446` builds
   `model_config` from `cfg["stage1_model"]` / `cfg["stage2_model"]`, and `cfg` always
   originates from `merge_cli_overrides(DEFAULT_CONFIG, ...)`, so every key is present
   and no fallback fires.
2. The legacy path also populates the keys explicitly.
   `_migrate_legacy_model_config` (`network.py:1446`) hard-codes
   `"decoder": "one_shot"` and `"particle_type": {"target": "physical"}` into the dict it
   returns, so a v0.2 checkpoint does not rely on the fallbacks either.
3. The tests build their configs by deep-copying `DEFAULT_CONFIG` (e.g.
   `tests/test_network.py:611` `_minimal_model_config`), so they never exercise the
   fallback branch.

That is precisely what makes this dangerous rather than harmless. The fallbacks are
**unreachable by every current caller, untested, and wrong.** The first caller that
hand-builds a partial `model_config` — a new test, a notebook, a debugging script, a
future `giant train --stage2-only` path — silently gets the v0.2 architecture while the
config documentation promises the v0.3 one.

### How to verify

```bash
uv run python - <<'EOF'
import copy
from giant import config as gc
from giant.model.network import build_models

full = copy.deepcopy(gc.DEFAULT_CONFIG)
print("DEFAULT_CONFIG stage2 decoder      :", full["stage2_model"]["decoder"])
print("DEFAULT_CONFIG particle_type.target:", full["stage2_model"]["particle_type"]["target"])

partial = {
    "pdg_vocab": 3, "mat_vocab": 2,
    "conditioning": full["conditioning"],
    "stage1_model": {"hidden_dim": 8, "n_res_blocks": 1},
    "stage2_model": {"hidden_dim": 8, "n_res_blocks": 1, "k_max": 3},
}
print("build_models fallback stage2 class :", type(build_models(partial)["stage2"]).__name__)
EOF
```

Observed output on `55332db`:

```
DEFAULT_CONFIG stage2 decoder      : autoregressive
DEFAULT_CONFIG particle_type.target: onehot
build_models fallback stage2 class : Stage2OneShot
```

### Recommended fix

Make the config **typed at the boundary**, so the defaults exist exactly once and the
type checker can see them.

1. Define a frozen dataclass per config block in `giant/config.py`:
   `ConditioningAxisConfig`, `ConditioningConfig`, `RouterConfig`, `WganConfig`,
   `Stage1ModelConfig`, `Stage2ModelConfig`, `TrainConfig`, `GiantConfig`. The dataclass
   field defaults become the *only* declaration of each default. Move the existing
   `DEFAULT_CONFIG` comments onto the fields — they are the most valuable thing in that
   file and must not be lost.
2. Give each a `from_dict(d)` classmethod that constructs from the merged TOML dict and
   **rejects unknown keys** (that is Issue 2, and the same constructor solves both).
3. Change `build_models`, `build_critics` and `StageSpec.from_config` to take the
   dataclasses instead of dicts. Every `.get(key, default)` becomes `cfg.key`. The 132
   duplicate defaults disappear by construction.
4. Keep `DEFAULT_CONFIG` as a derived artifact if anything still needs the dict shape
   (`save_config`, `default_out_dir_name`, the setup-cache sidecar) — generate it from
   the dataclasses via `dataclasses.asdict`, do not maintain it by hand.

Precedent exists in this codebase and should be followed:
`@dataclass(frozen=True) StageSpec` (`trainers.py:79`) and `PlotSpec`
(`analysis/catalog.py:65`) already do exactly this, well.

### Interim fix if the full change is too large

If the dataclass migration is deferred, at minimum add a test that asserts the inline
fallbacks agree with `DEFAULT_CONFIG`, so the drift is caught:

```python
def test_build_models_fallbacks_match_default_config():
    """Every `.get(key, default)` in build_models must use the same default
    DEFAULT_CONFIG declares — see issues.md Issue 1."""
```

and fix the two divergent values. Decide deliberately which way: if the fallbacks are
meant to encode *v0.2 legacy* behaviour, they are redundant (the legacy migration at
`network.py:1446` already sets both keys explicitly) and should be **deleted**, letting
a missing key raise `KeyError` instead of silently choosing an architecture.

### Scope guard

Do **not** change any actual default *value* as part of this refactor. The goal is to
make the two declarations agree and then have only one; changing what the model does is
a separate, physics-relevant decision that belongs in its own commit with its own
retraining. When resolving the two divergences above, the correct target is whatever
`DEFAULT_CONFIG` says (`autoregressive` / `onehot`), because that is the documented
v0.3.0 intent per `CLAUDE.md`.

---

## Issue 2 — No unknown-key validation: a typo in `config.toml` silently trains the wrong model

> **Status: Fixed.** `giant/config.py` now has `validate_config_keys(cfg)`, which walks a
> merged config dict recursively against `DEFAULT_CONFIG`'s tree (itself generated from
> the `GiantConfig` dataclasses added for Issue 1) and raises `ValueError` on any key not
> present there, with a `difflib`-based "did you mean" suggestion. `[meta]` is skipped
> unconditionally, and `stage{1,2}_model.router`'s `axis{i}_{field}` composed-router keys
> and runtime-seeded `centers_init` are allowed through explicitly. `merge_cli_overrides`
> calls it right before returning, so all three real entry points (`giant train`,
> `giant new-run`, `scripts/warm_setup_cache.py`) are covered automatically; checkpoint
> `model_config` loading (a separate code path, `network._migrate_legacy_model_config`)
> is untouched, so old checkpoints keep loading regardless of schema drift. This is the
> "standalone fallback" option from the recommended fix below rather than the
> dataclass-`from_dict`-rejects-unknown-keys option, specifically to keep the check scoped
> to the config.toml/CLI-overrides path without touching `validate_config`'s many direct
> unit-test callers or `run_train_job`'s redundant internal `validate_config` call.
> Regression tests added in `tests/test_config.py` cover top-level and nested typos, the
> axis/`centers_init` allowances, `[meta]` skipping, and both the file and CLI-overrides
> paths, plus the two real `configs/*.toml` fixtures that don't already fail
> `migrate_config` for unrelated reasons. Everything below this point describes the
> pre-fix state and is kept for historical context.

**Severity: High. Effort: Small. Risk if unfixed: wasted GPU-days on a run that did not
use the setting you thought it did.**

**Location:** `giant/config.py:456` (`_deep_merge`), `giant/config.py:628`
(`merge_cli_overrides`), `giant/config.py:661` (`validate_config`).

### What the code does today

`merge_cli_overrides` resolves configuration as `DEFAULT_CONFIG → TOML file → CLI
overrides`, deep-merging at each step. `_deep_merge` accepts any key:

```python
def _deep_merge(base: dict, override: dict) -> dict:
    result = dict(base)
    for k, v in override.items():
        if isinstance(v, dict) and isinstance(result.get(k), dict):
            result[k] = _deep_merge(result[k], v)
        else:
            result[k] = v          # <-- unknown k is accepted and stored
    return result
```

`migrate_config` explicitly preserves unrecognised sections
(`giant/config.py:612-615`: *"Anything else in the original dict (unrecognized top-level
sections) carries through untouched rather than being silently dropped"*).

`validate_config` (`giant/config.py:661`) exists and is *good* — it catches cross-block
contradictions with genuinely excellent, actionable error messages (e.g. "`router.type =
'pdg'` builds its own training-vocab-scoped embedding, incompatible with
`conditioning.particle.type = 'physical'`"). But by design it only inspects keys it knows
about. It cannot detect a key that should not exist.

### Why it's a problem

A config file containing

```toml
[stage1_model]
n_res_block = 12        # typo: should be n_res_blocks
```

merges cleanly, validates cleanly, trains to completion, and reports success — having
built a 6-block model. The stray `n_res_block` key is faithfully written back into the
run's saved `config.toml` by `save_config`, so the artifact of record also claims the
typo was a real setting. Nothing anywhere in the pipeline will ever say otherwise.

The blast radius scales with the config surface. v0.3.0 introduced a large nested schema
(`conditioning` × `stage1_model` × `stage2_model` × per-generator sub-tables × router
sub-tables), which is a great deal of surface for typos, and the v0.3.0 workflow is
explicitly *"a sequence of architecture comparisons"* (`config.py`, the `wandb` default
comment). A comparison in which one arm silently ignored its distinguishing setting is
worse than no comparison — it produces a confident, wrong conclusion.

There is a partial mitigation already: `warn_if_git_hash_mismatch` and
`warn_if_checkpoint_config_mismatch` warn about *provenance* drift. Neither looks at key
names.

### How to verify

```bash
uv run python - <<'EOF'
import tempfile, pathlib
from giant import config as gc

p = pathlib.Path(tempfile.mkdtemp()) / "config.toml"
p.write_text("[meta]\nconfig_version = 3\n\n[stage1_model]\nn_res_block = 12\n")
cfg = gc.merge_cli_overrides(gc.DEFAULT_CONFIG, p, {})
gc.validate_config(cfg)                      # passes
print("typo key survived :", cfg["stage1_model"]["n_res_block"])
print("real key unchanged:", cfg["stage1_model"]["n_res_blocks"])
EOF
```

### Recommended fix

Add strict key validation against the known schema. Two options, in order of preference:

1. **Preferred — fold into Issue 1.** If each config block becomes a dataclass with a
   `from_dict` that raises on unknown keys, this issue is solved for free and cannot
   regress. Raise a `ValueError` in the established house style, naming the offending
   key, its section, and the closest valid key by edit distance
   (`difflib.get_close_matches`) — that last part matters, because "unknown key
   `n_res_block`; did you mean `n_res_blocks`?" is the message that actually saves the
   run.

2. **Standalone fallback.** Add `validate_config_keys(cfg)` to `giant/config.py` that
   walks `cfg` alongside `DEFAULT_CONFIG` and raises on any key not present in the
   defaults tree. Call it from `merge_cli_overrides` right before returning, so every
   entry point gets it.

Either way, the following need explicit allowances, since they legitimately carry keys
not in `DEFAULT_CONFIG`:

- **`meta`** — holds `config_version`, `git_hash`, and run metadata from
  `build_run_meta`, none of which are in the defaults tree.
- **Composed-router axis keys** — `stage{1,2}_model.router.axis{i}_{field}` are
  deliberately dynamic and deliberately absent from `DEFAULT_CONFIG` (see the comment at
  the end of the `router` block in `config.py`, and `_parse_composed_axes` in
  `network.py:499`). Validate these against the `axis(\d+)_(\w+)` pattern plus the known
  per-axis field names rather than against a fixed key list.
- **`stage2_model.n_sec.legacy_owner`** — injected by `_migrate_legacy_model_config`, not
  a user-facing key. It appears in checkpoint `model_config` dicts, not in `config.toml`,
  so it should not reach this validator; confirm that before assuming.

### Scope guard

Strict validation must apply to the **`config.toml` path only**, not to loading old
checkpoints. A v0.2 checkpoint's `model_config` is migrated by a different function
(`network._migrate_legacy_model_config`, see Issue 6) and must keep loading. Adding
strictness that rejects historical checkpoints would break `giant predict` / `giant
rollout` against every model trained so far.

---

## Issue 3 — `cli.py:train()` is a 58-parameter, 510-line fat controller

> **Status: Fixed, commit `2bfb1ab` on `v0.3.0-stage2-autoregressive`.** `giant/config.py`
> now declares `FlagSpec` (a frozen dataclass: `name`, `paths` — one or more dotted config
> paths, `precedence`) and a `FLAG_SPECS` table covering every flag `train`/`new-run` map
> into the config tree, plus `overrides_from_flags(values: dict[str, object]) -> dict`,
> which applies specs in ascending precedence order (so a more-specific flag overwrites a
> shared/shorthand one written earlier at the same path) — one mechanism replacing the
> three different ad hoc "more specific wins" patterns identified below
> (`.update()`-call-order, `setdefault(...)[...] =` overwrite-order, and
> `{**shared, **specific}` merge). `train()`'s ~140-line override-building block
> (`cli.py:656-790` pre-fix) is now a single flat `flag_values` dict (mostly enum `.value`
> unwrapping) plus one call to `overrides_from_flags`; `new_run()`'s near-verbatim copy
> (`cli.py:915-983` pre-fix) collapsed the same way, reusing the identical table. Router
> overrides (`_router_cli_overrides`, unchanged, still shared by both commands) feed into
> the table as a single pre-aggregated `router_config` entry mapped only to
> `stage1_model.router` — the stage1-only asymmetry noted below is preserved exactly, with
> a regression test. No flag was added, removed, or renamed, and no precedence semantics
> changed: `giant train --help`/`giant new-run --help` are byte-identical before and after,
> verified by diffing both. Everything below this point describes the pre-fix state and is
> kept for historical context.

**Severity: High. Effort: Medium.**

**Location:** `giant/cli.py:329-839`.

### What the code does today

Measured on `55332db`:

| Command | Parameters | Lines |
|---|---|---|
| `train` | **58** | **510** |
| `new_run` | 31 | 188 |
| `predict` | 9 | 387 |
| `rollout` | 14 | 237 |

Of `train`'s 510 lines, roughly 430 are the Typer signature (one `Annotated[...]` block
per flag) and roughly 80 are hand-written translation from flags into the nested
overrides dict. That translation follows the same shape five times over — for `train`,
`stage1_model`, `stage2_model`, `conditioning`, and the WGAN sub-tables:

```python
cli_stage1_model: dict[str, object] = {
    k: v for k, v in {
        "hidden_dim": hidden_dim,
        "n_res_blocks": n_blocks,
        "dropout": dropout,
    }.items() if v is not None
}
cli_stage1_model.update(
    {k: v for k, v in {
        "hidden_dim": stage1_hidden_dim,
        "n_res_blocks": stage1_n_res_blocks,
        "dropout": stage1_dropout,
    }.items() if v is not None}
)
```

On top of that sits genuinely intricate precedence logic, correctly implemented but
expressed imperatively:

- `--mode` sets `generator` on **both** stages; `--stage1-generator` / `--stage2-generator`
  then override a single stage (`cli.py:713-720`).
- `--hidden-dim` / `--n-blocks` / `--dropout` are stage-1-only backward-compatible
  shorthands; `--stage1-*` wins when both are given (`cli.py:676-700`).
- `--n-critic` / `--gp-weight` / `--noise-dim` / `--critic-lr` fan out to
  `stage{1,2}_model.wgan.*`, with `--stage{1,2}-*` variants overriding per stage — built
  by merging a shared dict under a stage-specific one (`cli.py:722-753`).
- `--emb-dim` and `--conditioning` each set **two** places
  (`conditioning.particle.*` and `conditioning.material.*`) (`cli.py:702-710`).

### Why it's a problem

1. **The mapping is data, written as code.** Flag → dotted config path → fan-out rule is
   a table. Written as 80 lines of copy-adapted dict comprehensions, correctness has to
   be re-verified by reading, every time a flag is added. `giant/config.py:446` already
   provides `_set_path(d, "stage1_model.wgan.n_critic", v)` — the primitive a table-driven
   version needs.
2. **Adding a flag touches three places at once** (signature, the right comprehension,
   the right precedence rule) with no mechanism forcing them to stay consistent. v0.3.0's
   design explicitly anticipates more per-stage flags.
3. **It is the least-tested code in the repo** (Issue 4), because unit-testing a
   58-parameter Typer command means going through `CliRunner` with argv strings.
4. It obscures the ~15 lines that actually do something (`cli.py:825-839`): resolve the
   out-dir, echo two lines, call `run_train_job`.

### Recommended fix

Extract the mapping into a pure, table-driven function in `giant/config.py` (not
`cli.py`), so it is importable and directly testable without Typer:

```python
@dataclass(frozen=True)
class FlagSpec:
    """One CLI flag's mapping into the config tree."""
    name: str                    # "stage1_hidden_dim"
    paths: tuple[str, ...]       # ("stage1_model.hidden_dim",) — >1 means fan-out
    precedence: int = 0          # higher wins; per-stage flags outrank shorthands

FLAG_SPECS: tuple[FlagSpec, ...] = (
    FlagSpec("hidden_dim",        ("stage1_model.hidden_dim",), precedence=0),
    FlagSpec("stage1_hidden_dim", ("stage1_model.hidden_dim",), precedence=1),
    FlagSpec("mode",              ("stage1_model.generator", "stage2_model.generator")),
    FlagSpec("n_critic",          ("stage1_model.wgan.n_critic", "stage2_model.wgan.n_critic")),
    FlagSpec("emb_dim",           ("conditioning.particle.emb_dim", "conditioning.material.emb_dim")),
    ...
)

def overrides_from_flags(values: dict[str, object]) -> dict:
    """Build the nested overrides dict from {flag_name: value}, dropping
    None (= flag not given) and applying precedence."""
```

`train()` then becomes: parse `--batch-size auto`, collect `locals()`-style flag values,
call `overrides_from_flags`, `merge_cli_overrides`, `validate_config`, resolve out-dir,
call `run_train_job`. The 430-line Typer signature stays — that is irreducible, it *is*
the user interface — but it becomes the only bulk in the function.

The precedence table also becomes self-documenting, which today requires reading three
separate inline comment blocks to reconstruct.

### Scope guard

- **Do not remove or rename any flag.** The backward-compatible shorthands
  (`--hidden-dim`, `--n-blocks`, `--dropout`, `--mode`) exist because they predate the
  per-stage flags and are in people's shell history and job scripts. Their surprising
  stage-1-only scoping is documented behaviour, not a bug to fix.
- **Do not change precedence semantics.** Reproduce the current rules exactly, then
  test them (Issue 4). Any intentional change is a separate commit.
- Preserve the explanatory comments at `cli.py:676-680` and `cli.py:713-719` — they
  record *why* the precedence is the way it is.

---

## Issue 4 — `cli.py` sits at 35.8 % coverage and holds untested override-precedence logic

> **Status: Fixed for the override-precedence logic (Issue 3's scope); the
> `predict`/`rollout` inference-bootstrap portion described below is still open — that is
> Issue 5, deliberately not attempted here.** Commit `2bfb1ab` on
> `v0.3.0-stage2-autoregressive` adds direct, `CliRunner`-free unit tests for
> `overrides_from_flags` in `tests/test_config.py` — one per precedence rule, including
> several with previously **zero** coverage: both legs of `--stage{1,2}-generator`
> overriding `--mode` (only the stage1 leg had a test before), all three stage1
> shorthand-vs-`--stage1-*` pairs (previously only `hidden_dim`), all four WGAN knobs'
> shared-vs-per-stage precedence (previously only `n_critic`/`gp_weight`), the
> `--emb-dim`/`--conditioning` dual-axis fan-out (previously untested via `train` at all),
> and an explicit regression test that `router_config` only ever writes
> `stage1_model.router`. `tests/test_cli_train_overrides.py` keeps its original 4
> `CliRunner` smoke tests unmodified, plus 4 new ones covering the `--batch-size auto`
> parse-error and success paths and all three `out_dir` resolution branches
> (`--out`/`--resume`/default) — previously entirely uncovered. The percentage barely moves
> (36 % → 35 %), because the extraction *deleted* more statements from `cli.py` (539→478)
> than the new tests cover elsewhere in the file, but covered statements rose in absolute
> terms (157→169) and every override-precedence line the original issue called out by
> number is now covered. Coverage of `cli.py`'s other listed gap — the `predict`/
> `rollout` inference bootstrap (`cli.py:1104-1423`/`1528-1702` in the pre-fix numbering) —
> is unchanged, since fixing that requires the `load_for_inference` extraction described in
> Issue 5, which is explicitly out of scope for this fix (done separately, if at all).
> Everything below this point describes the pre-fix state and is kept for historical
> context.

**Severity: High. Effort: Medium.**

**Location:** `giant/cli.py` (1865 lines).

### What the code does today

Per-module line coverage from `coverage.xml`:

| Module | Coverage |
|---|---|
| `rollout.py` | 99.6 % |
| `analysis/catalog.py` | 99.6 % |
| `training/trainers.py` | 99.5 % |
| `model/network.py` | 98.6 % |
| `pipeline.py` | 97.1 % |
| `data/transforms.py` | 96.2 % |
| `config.py` | 91.8 % |
| **`cli.py`** | **35.8 %** |

Project total is 81.8 %. `cli.py` is not merely the lowest — it is a **60-point outlier**
against an otherwise uniformly high standard, and it is simultaneously the largest module
in the repo.

Existing CLI tests (`tests/test_cli_train_overrides.py`, `tests/test_cli_new_run.py`,
`tests/test_cli_predict.py`) are the right idea and well-written; they simply cover a
fraction of the surface.

### Why it's a problem

The uncovered code is not boilerplate. It includes:

- The entire flag → config override translation and its precedence rules (Issue 3). Every
  "`--stage2-hidden-dim` beats `--hidden-dim`", "`--stage2-generator` beats `--mode`",
  "`--n-critic` fans out to both stages unless `--stage1-n-critic` is given" rule is
  currently **unverified by any test**.
- The complete inference bootstrap in `predict` and `rollout` — checkpoint loading,
  normalizer reconstruction, conditioning-axis resolution, top-N map wiring (Issue 5).
  This is the code where a divergence between the two commands produces silent
  train/inference skew rather than a crash.
- The `--batch-size auto` parsing and estimation paths, and the six distinct
  "retrain with the current code" checkpoint-compatibility guards.

A failure here is expensive in a way a unit-test failure is not: a mis-scoped flag means
a multi-hour GPU run trains a model that differs from the one the experiment log claims.
Given `CLAUDE.md` describes v0.3.0 as a sequence of architecture comparisons, this is the
worst possible place for silent divergence.

### How to verify

```bash
uv run pytest -q --cov=giant --cov-report=term-missing:skip-covered 2>&1 | grep "cli.py"
```

### Recommended fix

Coverage here is a *consequence* of Issue 3, not an independent goal. Sequence the work:

1. **First, extract the testable logic** (Issue 3): `overrides_from_flags` in
   `config.py`, and `load_for_inference` in a new `giant/checkpoint_io.py` (Issue 5).
2. **Then test the extracted functions directly** — no Typer, no `CliRunner`, no
   filesystem. One test per precedence rule:

   ```python
   def test_stage1_hidden_dim_beats_hidden_dim_shorthand(): ...
   def test_mode_sets_both_stage_generators(): ...
   def test_stage2_generator_overrides_mode_for_stage2_only(): ...
   def test_n_critic_fans_out_to_both_wgan_subtables(): ...
   def test_stage1_n_critic_overrides_shared_n_critic_for_stage1_only(): ...
   def test_emb_dim_sets_both_conditioning_axes(): ...
   ```

3. **Keep a thin layer of `CliRunner` smoke tests** over `train --help`, `predict`,
   `rollout` to catch signature/wiring breakage, but do not try to reach high coverage
   through the CLI surface — that is slow and brittle.

Target: get `cli.py` to roughly the repo norm by *moving code out of it*, not by writing
elaborate CLI-invocation tests. A `cli.py` that is genuinely just argument declaration
plus delegation can sit at modest coverage without concern, because there will be
nothing in it left to get wrong.

### Scope guard

Do not add `# pragma: no cover` to close the gap. The gap is a real signal and the fix is
extraction.

---

## Issue 5 — The inference bootstrap is duplicated verbatim between `predict` and `rollout`

> **Status: Fixed.** `giant/checkpoint_io.py` now holds a single
> `load_for_inference(checkpoint, device, command_name, weights="raw", require_stage2=True)`
> plus an `InferenceContext` dataclass (stage1/stage2 models, all three normalizers, both
> vocab maps, both top-N maps, both conditioning axes, `k_max`, both stages' ddpm step
> counts, `other_policy`, the raw `model_config`, and `epoch`/`best_val_loss` for
> `rollout`'s YAML sidecar) — exactly the design this issue proposed, verified against the
> current (not `55332db`-era) code before writing it. `predict`/`rollout` in `cli.py`
> each shrink to one `try/except CheckpointCompatibilityError` call plus a block of
> `ctx.<field>` unpacks; `cli.py` lost `_conditioning_axes`, `_stage_cfg`, `_ddpm_steps`,
> `_particle_type_other_policy`, `_load_pdg_topn_map`/`_load_mat_topn_map`, and
> `_load_model_weights` entirely (net ~180 lines off `cli.py`). The independent third copy
> in `giant/analysis/router_gating.py` was deleted in favour of a lazy
> `from giant.checkpoint_io import conditioning_axes` inside `load_router`'s existing
> lazy-import block, preserving that module's "polars/numpy only at module scope"
> contract (`checkpoint_io.py` imports torch eagerly, so it must never be imported at
> `router_gating.py` module scope). Two guard orderings from the two commands' drifted
> copies were consolidated into one (`warn_if_checkpoint_config_mismatch` now always
> runs right after the existence guards, and `predict`'s `--batch-size auto` estimate now
> reads `ctx.model_config` after the checkpoint loads rather than before) — both are
> console-output-order changes only, no error text or model behavior changed, confirmed by
> diffing `giant predict --help`/`giant rollout --help` byte-for-byte before and after (no
> flag touched) and re-reading both rewritten command bodies field-by-field against the
> original. `require_stage2` exists as a real parameter, exercised by a new test, even
> though both current callers pass the default `True`. New `tests/test_checkpoint_io.py`
> (17 tests: happy path, every guard individually with exact message-text assertions, the
> `require_stage2=False`/inactive-stage2 path, `conditioning_axes`/`stage_cfg` directly)
> plus thin `CliRunner` smoke tests in `tests/test_cli_predict.py` and the new
> `tests/test_cli_rollout.py` confirming `CheckpointCompatibilityError` actually surfaces
> as `typer.Exit(1)` through the CLI — previously this entire code path had zero test
> coverage. `uv run pytest -q` (803 passed, up from 784), `ruff check`, `ruff format
> --check`, and `ty check` all clean. Everything below this point describes the pre-fix
> state and is kept for historical context.

**Severity: High. Effort: Small. Risk if unfixed: silent train/inference skew.**

**Location:** `giant/cli.py:1121-1201` (`predict`), `giant/cli.py:1534-1593` (`rollout`),
plus a third partial copy at `giant/analysis/router_gating.py:86-120`.

### What the code does today

`predict` and `rollout` each contain ~65 lines that are near-identical line-for-line:

1. `ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)`
2. Guard: no `model_config` → "retrain with the current code" → `Exit(1)`
3. Guard: no `sec_decoder` → same
4. Guard: no `normalizer.sec_phys` → same
5. `_conditioning_axes(model_cfg)`
6. `_load_pdg_topn_map(ckpt)` / `_load_mat_topn_map(ckpt)`
7. Guard: `particle_conditioning == "onehot"` but no `pdg_topn_map` → `Exit(1)`
8. Guard: `material_conditioning == "onehot"` but no `mat_topn_map` → `Exit(1)`
9. `pdg_map` / `mat_map` key-type coercion (`{int(k): v ...}` / `{str(k): v ...}`)
10. Three `Normalizer.from_dict(ckpt["normalizer"][...])` calls
11. `build_models(model_cfg)`, unpack `stage1` / `stage2`
12. Guard: either stage `None` → "needs both" → `Exit(1)`
13. `_load_model_weights(...)`, `.to(device).eval()` on both
14. `typer.echo(f"loaded checkpoint: ...")`
15. `gconfig.warn_if_checkpoint_config_mismatch(checkpoint)`

Additionally, `_conditioning_axes` is **duplicated verbatim into a second module**.
`giant/analysis/router_gating.py:69` carries a byte-for-byte copy of
`giant/cli.py:74`, with the docstring openly acknowledging it: *"Mirrors
`giant.cli._conditioning_axes`."*

```bash
grep -rn "def _conditioning_axes" giant/
# giant/cli.py:74
# giant/analysis/router_gating.py:69
```

### Why it's a problem

This is not cosmetic duplication. This code decides **how input features are assembled at
inference time** — which conditioning mode is used, which top-N vocabulary maps are
applied, which normalizer statistics are restored. If `predict` and `rollout` ever
diverge on any of those, the result is not an exception. It is two commands producing
subtly different physics from the same checkpoint, with no error and no warning. That
class of bug is found by noticing that a plot looks wrong, weeks later.

The duplication is also **entirely untested** — it lives in the 35.8 %-covered region of
`cli.py` (Issue 4). And it is *already* growing: the third copy in `router_gating.py`
shows the pattern spreading into a package that is otherwise carefully isolated.

The self-aware "Mirrors `giant.cli._conditioning_axes`" comment is the tell. When a
developer documents a copy rather than removing it, the missing abstraction has been
identified but not yet built.

### Recommended fix

Create `giant/checkpoint_io.py` (name it whatever fits; the point is that it is **not**
`cli.py`) exposing one function and one result object:

```python
@dataclass(frozen=True)
class InferenceContext:
    """Everything needed to run a trained checkpoint forward, resolved once."""
    stage1: nn.Module
    stage2: nn.Module
    cond_norm: Normalizer
    tgt_norm: Normalizer
    sec_phys_norm: Normalizer
    pdg_map: dict[int, int]
    mat_map: dict[str, int]
    pdg_topn_map: TopNMap | None
    mat_topn_map: TopNMap | None
    particle_conditioning: str
    material_conditioning: str
    k_max: int
    stage1_ddpm_steps: int
    stage2_ddpm_steps: int
    other_policy: str
    model_config: dict

def load_for_inference(
    checkpoint: Path,
    device: torch.device,
    weights: str = "raw",
    require_stage2: bool = True,
) -> InferenceContext: ...
```

Raise a dedicated `CheckpointCompatibilityError` (carrying the current, genuinely helpful
message text) instead of calling `typer.echo` + `typer.Exit`; the CLI catches it and does
the echo/exit. That keeps the module free of Typer and makes it unit-testable.

`predict` and `rollout` each shrink by ~60 lines to a single call. `router_gating.py`
drops its `_conditioning_axes` copy and imports from the new module — note that
`router_gating` deliberately imports torch **lazily inside the function**
(`router_gating.py:88`) to keep the analysis workers' "polars/numpy only" contract; the
new module must be imported the same way there, and `checkpoint_io.py` must not be
imported at `analysis/` module scope.

### Scope guard

- Preserve every error message verbatim. They are specific and actionable
  ("checkpoint's `conditioning.particle.type='onehot'` but has no `pdg_topn_map` — retrain
  with the current code"), and users have seen them before.
- Preserve the `require_stage2` distinction: `giant predict` and `giant rollout` both need
  both stages today, but `stage1_model.active = false` / `stage2_model.active = false` are
  real config options (`DEFAULT_CONFIG`), so the parameter should exist rather than
  hard-coding the requirement.
- Do not fold `giant/analysis/router_gating.py`'s router-specific loading
  (`load_router`, which returns `None` for non-MoE checkpoints) into the shared function.
  It has different semantics — absence is a normal outcome there, not an error.

---

## Issue 6 — Two independent v0.2→v0.3 migration surfaces encode the same knowledge

**Severity: Medium. Effort: Medium.**

**Location:** `giant/config.py:514` (`migrate_config`) and
`giant/model/network.py:1446` (`_migrate_legacy_model_config`).

### What the code does today

v0.3.0 broke the config format: the v0.2 single `[train]` + `[model]` layout became
`[conditioning]` / `[stage1_model]` / `[stage2_model]` / `[train]`. That break has to be
absorbed in two different places, and today it is absorbed by two unrelated functions:

- **`config.migrate_config`** translates a v0.2 **`config.toml`** on load.
- **`network._migrate_legacy_model_config`** translates a v0.2 **checkpoint's
  `model_config` dict** on `build_models`.

The split is acknowledged as deferred work in `migrate_config`'s own docstring:

> *"Operates on the config.toml shape. A checkpoint's `model_config` dict (which
> additionally carries n_sec_head ownership and needs `network.build_models`'s
> cooperation) is a separate migration surface, deferred to the network.py refactor."*

Both functions independently encode the same translation facts:

| v0.2 concept | v0.3 destination | in `migrate_config` | in `_migrate_legacy_model_config` |
|---|---|---|---|
| `train.mode` / `model.mode` | `stage{1,2}_model.generator` | ✓ | ✓ |
| `model.emb_dim` | `conditioning.{particle,material}.emb_dim` | ✓ | ✓ |
| `model.conditioning` | `conditioning.{particle,material}.type` | ✓ | ✓ |
| `model.n_blocks` | `stage{1,2}_model.n_res_blocks` | ✓ | ✓ |
| `model.noise_dim` | `stage{1,2}_model.wgan.noise_dim` | ✓ | ✓ |
| conditioning MLP depth was always 2 | `conditioning.*.n_layers = 2` | ✓ | ✓ |
| `router.expert_hidden_dim` set → hard error | — | ✓ | ✓ (near-identical message) |
| n_sec head lived on stage 1 | `stage2_model.n_sec.legacy_owner` | ✗ | ✓ (only here) |

The `expert_hidden_dim` rejection is the clearest symptom: the same policy, the same
reasoning, two hand-maintained copies of a ~10-line error message
(`config.py:581-591` and `network.py:1475-1489`).

### Why it's a problem

1. **Drift.** A future correction to the v0.2 interpretation must be applied to both.
   Applying it to one produces a checkpoint that loads with different architecture than
   its own config file describes.
2. **`legacy_owner` leaks into the builder.** Because the checkpoint migration is
   downstream of the config migration rather than sharing it, `build_models` has to carry
   legacy-specific branching at four sites (`network.py:1612`, `1642`, `1663`, `1688`):

   ```python
   legacy_owner = (s2cfg.get("n_sec") or {}).get("legacy_owner")
   n_sec_head_k_max = s2cfg.get("k_max", K_MAX) if legacy_owner == "stage1" else None
   ...
   build_n_sec_head=legacy_owner != "stage1",
   ```

   A builder for the current architecture should not need to know where v0.2 put its
   `n_sec` head.
3. **The retention policy is undeclared.** `tests/legacy/network_v02_snapshot.py` is
   **1009 lines** of frozen v0.2 network code kept purely so migration can be tested
   against it. That is a substantial maintenance surface with no stated expiry.

### Recommended fix

1. **Extract the shared translation into one table.** A single
   `_V02_TO_V03_TRANSLATION` mapping (old dotted key → tuple of new dotted keys) plus one
   set of "v0.2 architectural facts with no config key" constants, consumed by both
   functions. `config.py:475-511` already has the beginnings of this
   (`_V02_TRAIN_PASSTHROUGH`, `_V02_MODEL_TO_BOTH_STAGES`,
   `_V02_TRAIN_TO_BOTH_STAGES_WGAN`) — extend that pattern and share it. Put the shared
   table in a neutral module (e.g. `giant/_migration.py`) so neither `config.py` nor
   `network.py` has to import the other.
2. **Move `legacy_owner` handling out of `build_models`.** The migration function should
   emit a `model_config` that `build_models` can consume without legacy branching — for
   example by emitting an explicit `stage2_model.n_sec.owner` key that the *current*
   schema also carries (with value `"stage2"` for new runs), so the builder reads one key
   with two valid values rather than a nullable legacy sentinel.
3. **Write down the retention policy.** Add to `CLAUDE.md`: which v0.2 checkpoints must
   remain loadable, until when, and what triggers dropping the shim and the 1009-line
   snapshot. Without that, nobody will ever feel authorised to delete it.

### Scope guard

- **v0.2 checkpoints must keep loading** until the policy above says otherwise. There are
  trained models on `/ceph` that predate v0.3.0 and analysis runs referencing them.
- Keep `tests/legacy/network_v02_snapshot.py` as a **frozen snapshot** — do not "clean it
  up", reformat it, or make it share code with current `network.py`. Its entire value is
  that it is an independent, unchanging record of what v0.2 did. If ruff/ty complain about
  it, exclude it rather than edit it (it is already excluded from coverage via
  `pyproject.toml`'s `omit = ["*/legacy/*"]`).

---

## Issue 7 — Positional tuple contracts between the data, model and training layers

**Severity: Medium. Effort: Small.**

**Location:** `giant/data/transforms.py:863-890` (`build_features`),
`giant/data/dataset.py:229-237` (`StreamingStepsDataset` yield),
`giant/training/trainers.py:574-582` (unpack), `giant/cli.py:1233` (unpack).

### What the code does today

`build_features` returns a **bare 9-tuple**:

```python
) -> tuple[
    np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray,
    np.ndarray, np.ndarray, Normalizer | None, Normalizer | None,
]:
    """Assemble (cond_cont, cond_cat, target_s1, n_sec, sec_cont, proc_idx,
    sec_type_idx) arrays. ..."""
```

The element *meanings* live only in the docstring. Call sites re-derive them positionally:

```python
# giant/cli.py:1233 — predict
cond_cont, cond_cat, target_raw, _, _, _, _, _, _ = build_features(...)
```

`StreamingStepsDataset` yields a **7-tuple** of tensors (`dataset.py:229`), unpacked
positionally in the trainer:

```python
# giant/training/trainers.py:574
(
    cond_cont, cond_cat, x1_s1, n_sec, sec_cont, proc_idx, sec_type_idx,
) = _batch_to_device(batch, device)
```

and again, with a *different arity*, in the WGAN path:

```python
# giant/training/trainers.py:737
cond_cont, cond_cat, n_sec, sec_cont, sec_type_idx = batch_tensors
```

`_batch_to_device` is typed `(batch: tuple, device) -> tuple` — no element information at
all.

### Why it's a problem

1. **Invisible to the type checker.** `ty check` passes on all of it. Reordering two
   `np.ndarray` elements — or inserting a new one in the middle — type-checks cleanly and
   produces a model trained on scrambled features. Several of these arrays are
   shape-compatible with each other (`n_sec`, `proc_idx` and `sec_type_idx` are all
   integer arrays), so a swap may not even produce a shape error at runtime.
2. **This is the hottest contract in the codebase.** It crosses three layer boundaries
   (`data` → `training`, `data` → `cli`) and is the mechanism by which the model receives
   physics. A silent corruption here is the most expensive possible failure: it does not
   crash, it just trains something wrong.
3. **The nine-underscore unpack** at `cli.py:1233` is unreadable and does not survive any
   change to the tuple.
4. **Two different batch arities** for the same conceptual batch (7 in `_compute`, 5 in
   `_stage2_real_and_fake`) means the reader must track which slice is in play.

### Recommended fix

Convert both to `NamedTuple`. Zero runtime cost, full `ty` visibility, tuple-unpacking
still works so migration is incremental:

```python
class StepFeatures(NamedTuple):
    """Output of build_features. Field order is load-bearing for existing
    positional unpacking — append only, never insert or reorder."""
    cond_cont: np.ndarray
    cond_cat: np.ndarray
    target_s1: np.ndarray
    n_sec: np.ndarray
    sec_cont: np.ndarray
    proc_idx: np.ndarray
    sec_type_idx: np.ndarray
    cond_normalizer: Normalizer | None
    target_normalizer: Normalizer | None

class StepBatch(NamedTuple):
    cond_cont: torch.Tensor
    cond_cat: torch.Tensor
    target_s1: torch.Tensor
    n_sec: torch.Tensor
    sec_cont: torch.Tensor
    proc_idx: torch.Tensor
    sec_type_idx: torch.Tensor
```

Then `cli.py:1233` becomes `feats = build_features(...)` / `feats.cond_cont`, and
`_batch_to_device` gets the real signature `(batch: StepBatch, device) -> StepBatch`.

Precedent exists in this codebase: `RolloutSummary` (`rollout.py:276`, a `TypedDict`) and
`SetupStageResult` (`pipeline.py:39`, a dataclass) already do this correctly.

Note that `StreamingStepsDataset` passes batches through a PyTorch `DataLoader` with
`batch_size=None` and `num_workers > 0`, so the batch type must survive worker-process
pickling and `pin_memory`. `NamedTuple` does — it is a plain tuple subclass, and
`pin_memory` recurses into tuples — but **verify this with a real multi-worker run**, not
just the test suite, since the tests may run single-process.

### Scope guard

Do not reorder any existing fields while converting. The whole point of the change is to
make future reordering safe; performing one during the conversion, when nothing yet
protects against it, is the single riskiest version of this change.

---

## Issue 8 — `network.py` is 1745 lines holding three distinct modules

**Severity: Medium. Effort: Small (mechanical).**

**Location:** `giant/model/network.py`.

### What the code does today

One file contains six unrelated concerns:

| Lines (approx.) | Concern |
|---|---|
| 28-213 | Primitives: `SinusoidalEmbedding`, `cat_col_layout`, `_make_axis_mlp`, `ConditionEncoder`, `ContextAdapter`, `ResBlock` |
| 215-575 | **The entire router subsystem**: `Router` base, `register_router`/`build_router` registry, `EnergyRouter`, `PdgRouter`, `ProcessRouter`, `ComposedRouter`, `build_composed_router`, `_parse_composed_axes`, `_check_router_conditioning_compat`, `_build_router_from_cfg` |
| 577-727 | Trunks: `ExpertTrunk`, `_route_forward`, `Trunk`, `MonolithicTrunk`, `RoutedTrunk`, `build_trunk` |
| 729-881 | **History encoders**: `HistoryEncoder`, `MarkovHistory`, `_CausalAttnBlock`, `AttentionHistory` |
| 883-1443 | Top-level models: `Stage1Model`, `Stage2OneShot`, `Stage2Autoregressive`, `CriticModel` |
| 1446-1568 | **Legacy migration**: `_migrate_legacy_model_config`, `migrate_legacy_state_dict` |
| 1570-1745 | Builders: `build_models`, `build_critics` |

The v0.3.0 refactor described in `CLAUDE.md` — *"`network.py` is refactored from ten
permutation classes into composable parts (encoder × trunk × objective)"* — clearly
landed and is a genuine improvement. The composition is visible and clean. The file
simply was not split to match.

### Why it's a problem

- The seams are already there in the class layout; the file just does not honour them. A
  reader looking for `AttentionHistory` has no reason to expect it 800 lines into a file
  whose name suggests "the network".
- The router subsystem alone is ~370 lines with its own registry, its own config parsing,
  and its own compatibility validation. It is a subsystem, not a section.
- Legacy migration (Issue 6) sitting in the same file as the current builders is exactly
  what lets `legacy_owner` bleed into `build_models`.
- It is a merge-conflict magnet on a repo with parallel feature branches
  (`condor-gpu-train-rollout`, `v0.3.0-stage2-autoregressive`).

### Recommended fix

Split along the existing seams into `giant/model/`:

```
giant/model/
    layers.py      # SinusoidalEmbedding, ResBlock, ContextAdapter, _make_axis_mlp
    encoders.py    # ConditionEncoder, cat_col_layout
    routers.py     # Router base + registry + all 4 router types + config parsing
    trunks.py      # Trunk, MonolithicTrunk, RoutedTrunk, ExpertTrunk, build_trunk
    history.py     # HistoryEncoder, MarkovHistory, AttentionHistory, _CausalAttnBlock
    models.py      # Stage1Model, Stage2OneShot, Stage2Autoregressive, CriticModel
    builders.py    # build_models, build_critics
    _legacy.py     # _migrate_legacy_model_config, migrate_legacy_state_dict
    network.py     # re-export shim: `from giant.model.layers import *` etc.
```

Keep `network.py` as a **re-export shim** so no import site outside `giant/model/` has to
change in the same commit. `network.py` is imported by `cli.py`, `sample.py`,
`trainers.py`, `pipeline.py`, `analysis/router_gating.py`, and heavily by the test suite —
`grep -rn "from giant.model.network import" giant/ tests/ | wc -l` before starting, to
size the blast radius if you do decide to update call sites (21 sites on `55332db`).

This is a pure file-move refactor: no logic changes, and the 725-test suite plus `ty
check` is a strong safety net for it.

### Scope guard

Do it as its own commit, containing **only** moves and imports. Do not combine with Issue 1
or Issue 6, both of which change behaviour in this file — a diff mixing moves with logic
changes is effectively unreviewable.

---

## Issue 9 — `scripts` is published as a top-level distribution package

**Severity: Medium. Effort: Small.**

**Location:** `pyproject.toml`:

```toml
[project.scripts]
giant = "giant.cli:app"
dwarf = "scripts.dwarf:app"

[tool.hatch.build.targets.wheel]
packages = ["giant", "scripts"]
```

### What the code does today

The repo's tooling CLI (`dwarf`) lives in a directory called `scripts/`, which is declared
as a wheel package and referenced by the console-script entry point as `scripts.dwarf:app`.
Installing `giant` therefore creates a top-level importable module named **`scripts`** in
`site-packages`.

`scripts/` holds ten real modules: `dwarf.py`, `bump_dataset_version.py`,
`create_root_files.py`, `geometry_oracle.py`, `hparam_scan.py`, `migrate_geant_steps.py`,
`profile_analysis_costs.py`, `steps_to_parquet.py`, `steps_to_parquet_parallel.py`,
`warm_setup_cache.py`.

### Why it's a problem

`scripts` is one of the most generic names possible in the Python ecosystem. Consequences:

1. **Collision.** Any other installed distribution that also ships a top-level `scripts`
   package silently shadows or is shadowed by this one, depending on `sys.path` order.
   The failure mode is an `ImportError` or — worse — importing someone else's `dwarf`-less
   `scripts` and getting `AttributeError` at CLI startup.
2. **Environment-order fragility.** On the portal machines (`/work/lbogner`, shared with
   other users, per `CLAUDE.md`), a stray `scripts/` directory in the CWD shadows the
   installed package, because CWD precedes `site-packages` on `sys.path`. Running `dwarf`
   from a directory that happens to contain a `scripts/` folder can break in a confusing
   way.
3. **`scripts/` reads as "not part of the product"**, yet it is installed, has an entry
   point, is covered by `[tool.coverage.run] source`, and has a full test suite
   (`tests/test_dwarf.py` at 81.7 % coverage, plus `test_steps_to_parquet*.py`,
   `test_create_root_files.py`, `test_bump_dataset_version.py`). Its name misrepresents
   its status.

### Recommended fix

Move it under the `giant` namespace, where it cannot collide:

```
giant/tools/          # was scripts/
    dwarf.py
    ...
```

```toml
[project.scripts]
giant = "giant.cli:app"
dwarf = "giant.tools.dwarf:app"

[tool.hatch.build.targets.wheel]
packages = ["giant"]
```

Then update:

- `[tool.coverage.run] source = ["giant"]` (drop the now-redundant `"scripts"`).
- All `from scripts.X import Y` in `tests/` (`grep -rn "scripts\." tests/`).
- Any `python scripts/foo.py` invocations in docs, `README.md`, `CLAUDE.md`, HTCondor
  submit files, or shell history/job scripts on the portal machines.

The `dwarf` **command name does not change**, so anything invoking `dwarf ...` keeps
working — only Python-level imports and direct `python scripts/...` paths move.

### Scope guard

Check `giant/analysis/condor.py` and any generated HTCondor submit descriptions for
hard-coded `scripts/` paths before moving. Jobs submitted to the cluster may reference the
path as it exists on `/work` or `/ceph`, and a rename that lands mid-flight breaks queued
jobs. Search for the literal string `scripts/` across the repo, not just Python imports.

---

## Issue 10 — `torch.load(weights_only=False)`: checkpoints are arbitrary pickles

**Severity: Low (given the threat model). Effort: Medium.**

**Location:** `giant/cli.py:1122`, `giant/cli.py:1535`,
`giant/analysis/router_gating.py:93`, `giant/training/loop.py:153`.

### What the code does today

All four checkpoint loads pass `weights_only=False`:

```python
ckpt = torch.load(checkpoint, map_location="cpu", weights_only=False)
```

This is *necessary* today, not careless: the checkpoint dict carries non-tensor objects
beyond raw state dicts — `model_config`, `pdg_map` / `mat_map`, `pdg_topn_map` /
`mat_topn_map`, and the `normalizer` sub-dict, assembled in
`giant/training/checkpoint.py:19` (`build_checkpoint`) from the `extras` argument.
`weights_only=True` would reject them.

### Why it's worth recording

`weights_only=False` means loading a checkpoint executes arbitrary pickle opcodes. The
practical threat model here is mild — checkpoints are produced by this codebase and live
on `/ceph/lbogner` — but the portal machines and `/ceph` are explicitly **shared with
other users** (`CLAUDE.md`, "Compute environment"), and `giant predict --checkpoint <path>`
will happily load a path someone else wrote. It is also the kind of thing that becomes a
blocker later, if a model is ever shared outside the group or published alongside a paper.

Secondary practical cost: a pickled `model_config` cannot be inspected without importing
torch and unpickling. `dwarf status`-style tooling, or a human answering "what conditioning
mode was this trained with?", has to load the whole checkpoint.

### Recommended fix

Split the checkpoint into a tensor part and a JSON sidecar:

- `ckpt.pt` — state dicts only, loadable with `weights_only=True`.
- `ckpt.meta.json` — `model_config`, vocab maps, top-N maps, normalizer statistics,
  `epoch` / `global_step` / `best_val_loss`.

This has real secondary benefits: the sidecar is greppable and diffable, so comparing two
runs' architectures becomes `diff` rather than a Python session, and it composes well
with the config-provenance machinery already in `config.py`
(`warn_if_checkpoint_config_mismatch`, `git_hash`).

Note that `giant/data/setup_cache.py` already establishes the "JSON sidecar next to the
data" pattern for exactly this kind of non-tensor state — follow its conventions.

### Scope guard

This is a **format change**, so it needs a compatibility path: `load_for_inference`
(Issue 5) reads the sidecar if present and falls back to the pickled keys if not. Do not
attempt this before Issue 5 lands — with three separate copies of the loading code, a
format change means three separate compatibility paths. After Issue 5, it is one.

Given the mild threat model, this is correctly the **lowest-priority** item here. It is
recorded so the decision is deliberate rather than accidental.

---

## Issue 11 — Minor items

### 11a — `echo=print` threaded through the pipeline instead of `logging`

**Location:** `giant/pipeline.py:305` (`run_train_job(..., echo=print)`),
`giant/pipeline.py:86` (`run_setup_stage(..., echo=...)`); 52 `typer.echo` calls in
`cli.py`.

Passing the output function as a parameter is *better* than hard-coding `typer.echo` deep
in the pipeline — it keeps `pipeline.py` free of Typer and makes output capturable in
tests. But at this scale it has costs: no severity levels (the `--num-workers` quota
warning at `pipeline.py:329` and routine progress output are indistinguishable), no
timestamps on long training runs, no way to route to a file without wrapping, and the
parameter has to be threaded through every function that might print.

Suggested: standard `logging` with a Typer/rich handler configured once in `cli.py`.
Module-level `log = logging.getLogger(__name__)` replaces the threaded `echo`. Keep
`typer.echo` for genuine CLI output (the `device:` / `out_dir:` banners, error messages
before `Exit(1)`) — those are interface, not logs.

Low priority; the current approach is defensible. Worth doing opportunistically if
`pipeline.py` is being touched anyway.

### 11b — `giant.particles` imports from `giant.data.loader`

**Location:** `giant/particles.py:30`.

```python
if TYPE_CHECKING:
    from giant.data.loader import TopNMap
```

`giant/particles.py` is a physics-domain module (PDG code decoding, mass/charge lookup).
`giant/data/loader.py` is data I/O. A domain module depending on an I/O module is a
layering inversion.

Mitigating factors, which is why this is minor: the import is `TYPE_CHECKING`-guarded, so
there is **no runtime cycle**, and `ty check` passes. It is a design smell, not a defect.

Suggested: move `TopNMap` to a neutral location — `giant/constants.py` or a new
`giant/types.py` — so both `particles.py` and `data/loader.py` depend on it rather than on
each other. Do this only if `TopNMap` is being touched for other reasons; it is not worth
a standalone commit.

---

## Recommended sequence

The issues are interdependent. This order minimises rework:

1. **Issue 5** — extract `load_for_inference` into `giant/checkpoint_io.py`. Small,
   self-contained, removes the highest-risk duplication, and unblocks Issue 10.
2. **Issue 3** — extract the flag→config mapping table into `config.py`.
3. **Issue 4** — test the two extracted units directly. Issues 3 and 5 make this cheap;
   attempting it first means testing through `CliRunner`, which is slow and brittle.
4. **Issue 1 + Issue 2 together** — typed config dataclasses. One migration solves both,
   and doing them separately means touching the same 132 call sites twice.
5. **Issue 6** — unify the migration surfaces (easier once config is typed).
6. **Issue 8** — split `network.py`. Pure file moves; do it as an isolated commit, after
   the logic changes in Issues 1 and 6 have settled, to keep both diffs readable.
7. **Issue 7** — `NamedTuple` the feature/batch contracts.
8. **Issue 9** — move `scripts/` → `giant/tools/`. Independent of everything else; can be
   done at any point.
9. **Issue 10** — checkpoint format split, if the threat model or a publication makes it
   worthwhile.
10. **Issue 11** — opportunistically.

## General guidance for whoever picks this up

- **Run the full suite on every change.** `uv run pytest -q` takes 18 seconds. There is no
  reason to batch up unverified changes.
- **Run all three checks before committing:** `uv run ruff check .`,
  `uv run ruff format .`, `uv run ty check .`. CI enforces all three.
- **Do not change model behaviour while refactoring.** Several issues touch code that
  determines what network gets built (Issues 1, 6) or what features it receives (Issues 5,
  7). A refactor that also changes a default silently invalidates every prior benchmark in
  `/home/lars/knowledge-base/experiments/`. If a behaviour change is warranted, it is a
  separate commit with a separate message saying so.
- **Preserve the comments.** This codebase's docstrings explain reasoning, trade-offs, and
  known limitations to an unusually high standard. When code moves, the comments move with
  it. When code is deleted, check whether the comment records a decision that still needs
  recording elsewhere.
- **`giant/analysis/` is the template.** It is declarative where the rest of the codebase
  is imperative (`PlotSpec` registry), explicit about its own contracts (the
  `compute_partial` / `finalize` split, `chunkable=False`, the "polars/numpy only on
  workers" boundary), and honest about its one exception (`router_gating.py`'s module
  docstring explains precisely why it is allowed to import torch and why that is still
  safe). When deciding how a refactored `config.py` or `cli.py` should look, look there.