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>
AttentionHistory (giant/model/network.py) adds causal self-attention over
the emitted-secondary prefix as the alternative to MarkovHistory, with a
parallel forward() for training and an init_cache()/step() KV-cache path
for sample.py's per-slot AR inference loop, wired into
Stage2Autoregressive via history="attention".
giant/train.py adds _stage2_tf_prob and _assemble_stage2_ar_inputs_scheduled,
mixing ground-truth history with a detached sample_secondaries_ar self-sample
per slot so teacher_forcing="scheduled"/"never" close the train/inference gap
teacher_forcing="always" always avoided; wired into both stage-2 AR trainers.
config.py's validate_config no longer rejects these two previously
unimplemented schema values.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Replaces the Stage2Autoregressive stub with a real per-token secondary
decoder: MarkovHistory summarizes the previous secondary, remaining-energy
fraction and slot index round out the per-token conditioning, and the
existing Trunk/MonolithicTrunk/RoutedTrunk machinery is reused unchanged by
batching all K_MAX tokens together under teacher forcing (one parallel pass,
no new trunk code). build_models/build_critics wire it in; the WGAN critic
stays whole-sequence, so build_critics needs no AR-specific path.
train.py's FlowDDPMStageTrainer/WGANStageTrainer gain a decoder branch,
sharing optimizer/EMA/checkpoint machinery with the one-shot path.
_assemble_stage2_real is now defined in terms of the new unflattened
_assemble_stage2_ar_target helper, removing a near-duplicate branch.
Also lands the §11.4 differentiability validation-obligation instrumentation
(trunk-gradient norm from the particle-type slice vs. the continuous slices,
for generator=wgan + particle_type.target=onehot) via backward hooks in
_relax_onehot_type_slice, decoder-agnostic and surfaced as two new
metrics.csv columns.
This also fixes the standing regression where any config not explicitly
overriding decoder="one_shot" crashed at build_models, since
stage2_model.decoder defaults to "autoregressive" — confirmed by removing
tests/test_pipeline.py's now-stale override so the default config runs
end-to-end against real synthetic data.
Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com>
Builds the shared top-N-plus-other PDG/material maps (pooling both primary
and secondary occurrences for PDG, directly targeting the meeting's
species-collapse failure mode) and wires up conditioning.{particle,material}
= "onehot" plus stage2_model.particle_type.target in ("onehot", "embedding")
end-to-end: setup-cache persistence, Stage2OneShot's type_head (flow/ddpm)
vs. folded+ST-Gumbel-relaxed adversarial slice (wgan), and the corresponding
CE/MSE training losses. particle_type.target = "physical" stays byte-for-byte
unchanged, keeping the v0.2 migration shim's bit-identical guarantee intact.
giant predict/rollout fail loudly on a onehot/embedding checkpoint until
full decode support lands in step 6.
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