Files
giant/issues.md
T
lars 6a21c3b908
CI / Format (ruff format) (push) Successful in 28s
CI / Lint (ruff check) (push) Successful in 29s
CI / Sync project version with tag (push) Has been skipped
CI / Lint (ruff check) (pull_request) Successful in 41s
CI / Type check (ty) (push) Successful in 44s
CI / Format (ruff format) (pull_request) Successful in 31s
CI / Sync project version with tag (pull_request) Has been skipped
CI / Type check (ty) (pull_request) Successful in 34s
CI / Tests (push) Successful in 3m24s
CI / Tests (pull_request) Successful in 3m18s
Mark issues.md Issues 3 & 4 as fixed
Records what commit 2bfb1ab actually changed and its scope, matching the
status-blockquote convention already used for Issues 1 and 2.

Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
2026-08-12 15:05:33 +02:00

60 KiB
Raw Blame History

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.

  • pytest725 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 9699.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 Open
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):

# 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),
# 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

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

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:

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:

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

[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

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

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 keysstage{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:

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.

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:

@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

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

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:

    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

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_mapExit(1)
  8. Guard: material_conditioning == "onehot" but no mat_topn_mapExit(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."

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.

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:

@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):

    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.

  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:

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

# 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:

# 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:

# 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 (datatraining, datacli) 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.

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

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).

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:

[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.

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

giant/tools/          # was scripts/
    dwarf.py
    ...
[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:

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.

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

  • ckpt.pt — state dicts only, loadable with weights_only=True.
  • ckpt.meta.jsonmodel_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.

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.


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 7NamedTuple 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.