Works through docs/v0.3.0-followups.md item by item, closing the gap
between the design doc and the shipped v0.3.0-stage2-autoregressive code:
1. validate.py: 7-tuple batch unpacking, sample_stage1/sample_stage2
dispatch, stage-2 particle-type-class marginal.
2. Stage-prefixed --stage1-*/--stage2-* CLI flags for train/new-run.
3. Thread stage2_model.k_max through loader/transforms/dataset/pipeline/
train instead of the hardcoded K_MAX constant.
4. Mixed conditioning.particle.type / conditioning.material.type support
end-to-end (data pipeline + dwarf warm-cache).
5. conditioning.share_stages = true: one shared ConditionEncoder instance
across both stages.
6. stage2_model.generator = "ddpm" formally deferred into design doc §11.2
(was silently unimplemented).
7. giant predict/rollout: implement conditioning.*.type = "onehot" via the
checkpoint's saved pdg_topn_map/mat_topn_map.
8. network.py's checkpoint-path model_config migration now fails loudly on
non-zero legacy expert_hidden_dim/expert_n_blocks, matching config.py's
TOML-load path (§4.2).
9. validate_config now rejects stage2_model.n_sec.mode = "truth" for a
rollout-capable checkpoint (§9).
Also cleared all pre-existing `ty check` noise (44 -> 0 diagnostics),
mostly a test-helper dict-unpack pattern that made every unrelated
constructor keyword look like a type error.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
- giant/sample.py: fix every sampler's call convention against
Stage1Model/Stage2OneShot's actual forward signatures (was still
calling model(x, t, cond_cont, cond_cat) positionally); add
sample_secondaries_ar (free-running AR loop, unsnapped history feature)
and sample_stage1/sample_stage2/resolve_n_sec dispatch helpers that read
each stage's generator_kind/decoder off the model instance itself.
- giant/particles.py: decode_topn_class (argmax + other_policy) and
decode_embedding_nearest (L1-snap + distance) turn a secondary's
"onehot"/"embedding" type prediction into a concrete PDG.
- giant/rollout.py: decode_secondary_identity routes all three
particle_type.target values to real mass/charge; per-stage generator
dispatch (drops the single shared `mode` string, adds ddpm support);
L1DistCollector accumulates the §11.3 embedding-distance diagnostic.
- giant/cli.py: drop the onehot/embedding-target rejection gate (narrowed
to the still-unimplemented conditioning.particle/material.type=onehot
axis); fix the dead model_cfg.get("mode") bug in predict/rollout.
- giant/analysis/: new type_embedding_l1_distance PlotSpec, wired through
the rollout YAML sidecar (no live-model call needed, unlike
router_gating -- the histogram is already pre-aggregated at rollout
time).
- Un-xfail every test that was blocked on this step (test_rollout.py,
test_flow.py, test_wgan.py, test_phase2.py, test_router.py,
test_validate.py); add test_sample.py, test_type_embedding_distance.py.
Known follow-up: giant/validate.py still unpacks the training val-batch
as a stale 6-tuple and doesn't use the new per-stage dispatch, so
marginal validation during training degrades gracefully with a warning
rather than working -- not in this step's scope.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Decomposes the ten permutation classes in giant/model/network.py into
the reusable parts from docs/v0.3.0-design.md §5: ConditionEncoder (now
independently configurable per particle/material axis), ContextAdapter,
Trunk/MonolithicTrunk/RoutedTrunk/ExpertTrunk, and the stage classes
Stage1Model/Stage2OneShot/CriticModel (Stage2Autoregressive stubbed,
raises NotImplementedError until step 4/5). build_models/build_critics
now return a dict keyed by stage and accept the new nested config shape,
with routed WGAN reachable for the first time (the old --mode wgan
--router rejection is gone) and stage2_model.router.tie_to_stage1
sharing a literal Router instance.
A v0.2 checkpoint's flat model_config auto-migrates via
_migrate_legacy_model_config + migrate_legacy_state_dict, preserving the
n_sec_head's attachment to Stage1Model (legacy_owner="stage1", design
doc §4.1). tests/test_migration_v02_v03.py proves this bit-identical
against a frozen v0.2 snapshot (tests/legacy/network_v02_snapshot.py)
for both flow and wgan, both conditioning modes.
scripts/check_migration_v02_v03.py is the real-checkpoint counterpart
for a portal machine with /ceph access.
giant/model/schedule.py's flow-matching/DDPM loss helpers are updated
to the new model-call convention (t as a keyword). giant/sample.py,
giant/rollout.py, and giant/validate.py are not yet updated (deferred
to design doc step 6) — their exercising tests are marked xfail with
that reasoning rather than silently broken.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Replaces the independent log_delta_e/log_edep targets with 2 additive-log-ratio
coordinates over the deposit/secondary/post-energy simplex (fractions of pre_E
summing to 1), so edep + e_sec + post_E == pre_E holds by construction after
decoding (softmax) rather than being learned approximately. Requires e_sec
(secondary energy) as a new conditioning input and a steps_to_parquet.py pass
to derive it from child track first-step energies.
ruff removed unused imports across analysis.py and several test files.
ty caught a wrong dict[int, int] annotation on StreamingStepsDataset's
mat_map (materials are strings) and a real bug in steps_to_parquet.py
where --compression none passed None to polars' write_parquet, which
only accepts the literal "uncompressed". Also narrows a few
Optional-typed attributes (ddpm_schedule, Normalizer.mean/std) with
asserts and aligns __getitem__'s parameter name with torch's Dataset
base class.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
step_length already encodes |post_pos - pre_pos| by definition, so a raw
post_pos target would duplicate that magnitude and could drift inconsistent
with step_length during sampling. Instead add travel_dir, a unit vector
(local frame) giving only the direction of pre_pos->post_pos; post_pos is
reconstructed at inference as pre_pos + step_length * travel_dir, keeping
the two self-consistent. Target grows from 6D to 9D.
Co-Authored-By: Claude Sonnet 4.6 <noreply@anthropic.com>
lars