chore: update changelog for v0.3.18

Reviewed-on: #94

.bumpversion.toml

@@ -1,5 +1,5 @@
 [tool.bumpversion]
-current_version = "0.3.17"
+current_version = "0.3.18"
 parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
 serialize = ["{major}.{minor}.{patch}"]
 search = "{current_version}"

.gitea/workflows/ci.yml

@@ -173,6 +173,52 @@ jobs:
             git push origin "refs/tags/$TAG"
           fi
 
+  update-badges:
+    name: Update README badges (version, test count)
+    needs: [ruff-check, ruff-format, type-check, test, bump-version]
+    if: github.ref == 'refs/heads/master' && github.event_name == 'push'
+    runs-on: ubuntu-latest
+    container:
+      image: docker.gitea.com/runner-images:ubuntu-latest
+      volumes:
+        - /srv/act-runner-cache/uv:/uv-cache
+    steps:
+      # ref: master (not the triggering SHA) so this picks up whatever
+      # bump-version just pushed, rather than badging the pre-bump commit.
+      - uses: actions/checkout@v4
+        with:
+          token: ${{ secrets.CI_TOKEN }}
+          ref: master
+      - uses: astral-sh/setup-uv@v5
+        with:
+          enable-cache: false
+      - run: |
+          echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
+          echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
+      - run: uv sync --extra cpu --extra dev
+      - name: Compute version and test count
+        id: stats
+        run: |
+          VERSION=$(uv version --short)
+          TEST_COUNT=$(uv run pytest --collect-only -q 2>/dev/null | grep -oE '^[0-9]+ tests? collected' | grep -oE '^[0-9]+')
+          echo "version=$VERSION" >> "$GITHUB_OUTPUT"
+          echo "test_count=$TEST_COUNT" >> "$GITHUB_OUTPUT"
+      - name: Rewrite badge lines in README.md
+        run: |
+          sed -i -E "s|badge/version-[^-]+-informational|badge/version-${{ steps.stats.outputs.version }}-informational|" README.md
+          sed -i -E "s|badge/tests-[0-9]+%20passing-brightgreen|badge/tests-${{ steps.stats.outputs.test_count }}%20passing-brightgreen|" README.md
+      - name: Commit and push if the badges actually changed
+        run: |
+          git config user.name "gitea-actions"
+          git config user.email "actions@git.larsbogner.de"
+          git add README.md
+          if ! git diff --cached --quiet -- README.md; then
+            git commit -m "chore: update README badges (version ${{ steps.stats.outputs.version }}, ${{ steps.stats.outputs.test_count }} tests)"
+            git push origin HEAD:master
+          else
+            echo "Badges already up to date"
+          fi
+
   sync-version-on-tag:
     name: Sync project version with tag
     if: startsWith(github.ref, 'refs/tags/')

CHANGELOG.md

@@ -1,5 +1,13 @@
 # Changelog
 
+## [0.3.18] - 2026-09-03
+
+### Changed
+
+- Docs: rewrite README, keep version/test badges live via Gitea Actions
+
+- Perf: compact Stage-2 AR inference loop to active rows only
+
 ## [0.3.17] - 2026-09-02
 
 ### Changed

README.md

@@ -1,203 +1,295 @@
-# giant
+<div align="center">
 
-**G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep surrogate.
+# GIANT
 
-A conditional generative model that replaces the Geant4 step function: given a pre-step particle state it samples a post-step outcome — the primary's continuation plus its secondary particles — and autoregressively rolls that out into full showers. Trained entirely from parquet dumps of the miniCaloSim steps tree; no Geant4 runtime dependency.
+### **G**eant4 **I**nference via **A**utoregressive **N**eural s**T**ep surrogate
+
+A conditional generative model that replaces the Geant4 step function — sample a
+post-step outcome instead of simulating one, then roll that out into full
+calorimeter showers.
+
+[![python](https://img.shields.io/badge/python-3.12%2B-3776AB?logo=python&logoColor=white)](pyproject.toml)
+[![torch](https://img.shields.io/badge/torch-2.3.x-EE4C2C?logo=pytorch&logoColor=white)](pyproject.toml)
+[![version](https://img.shields.io/badge/version-0.3.17-informational)](CHANGELOG.md)
+[![tests](https://img.shields.io/badge/tests-1131%20passing-brightgreen)](tests/)
+[![CI](https://git.larsbogner.de/lars/giant/actions/workflows/ci.yml/badge.svg?branch=master)](https://git.larsbogner.de/lars/giant/actions)
+[![license](https://img.shields.io/badge/license-unlicensed-lightgrey)](#license)
+
+</div>
+
+---
+
+## The idea
+
+Geant4's step function is the innermost loop of detector simulation — for every
+particle, at every step, it stochastically samples where the particle goes next,
+how much energy it deposits, and what secondaries it spawns. GIANT learns that
+function instead of running it: given a pre-step particle state (position,
+energy, direction, particle species, material), a two-stage model samples a
+post-step outcome — including the variable-length list of secondaries — and
+autoregressively rolls that out into whole showers. It trains entirely from
+parquet dumps of a Geant4 steps tree; nothing downstream needs a Geant4 runtime.
+
+Two guarantees are architectural, not learned:
+
+- **Energy is conserved by construction.** Stage 1 decodes deposit / secondary /
+  post-step energy through a softmax simplex that sums to the pre-step energy
+  exactly; Stage 2's secondaries stick-break that same energy budget.
+- **No shower leaks across the train/val split.** Steps are split by `event_id`,
+  never by row, so correlated steps from the same shower can't appear on both
+  sides.
+
+## How a step becomes a shower
+
+```mermaid
+flowchart LR
+    A["pre-step state\nposition · energy · direction\nspecies · material"] --> B["ConditionEncoder\nphysical / embedding / onehot"]
+    B --> C["Stage 1\n9D post-step outcome"]
+    C --> D["Stage 2 (autoregressive)\nsecondaries, descending energy"]
+    D --> E["rollout step"]
+    E -->|"primary continues"| F["GeometryOracle\nposition to material, layer"]
+    E -->|"secondaries pushed"| G["track queue"]
+    F --> A
+    G --> A
+    E -->|"terminated"| H["deposited shower"]
+```
 
 ## Quick start
 
 ```bash
-uv sync --extra cpu               # install deps (CPU torch; use --extra cuda for GPU)
+uv sync --extra cpu                                  # install deps (CPU torch; --extra cuda for GPU)
-
-giant new-run --hidden-dim 512 --lr 3e-4          # scaffold config.toml + run dir
-giant model summary --config config.toml          # parameter counts + which config keys actually bite
-giant train path/to/steps.parquet                 # train (flow + wgan by default)
-giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt
 
+giant new-run --hidden-dim 512 --lr 3e-4              # scaffold config.toml + run dir
+giant train path/to/steps.parquet                     # train (flow stage 1 + wgan stage 2, by default)
 dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl   # needed for rollout
 giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl
+
+giant analyze prep rollout.yaml && giant analyze render <run_dir> --gallery   # rollout-vs-Geant4 diagnostics
 ```
 
-Every command takes `--help` for the full flag list, and `--config config.toml` for anything not exposed as a flag.
+Every command takes `--help` for its full flag list, and `--config config.toml`
+for anything not exposed as a flag.
 
-## Architecture
+---
 
-A **two-stage model**, checkpointed together. Either stage's outcome can be produced by one of three interchangeable generative objectives (`--stage1-generator`/`--stage2-generator`, or `--mode` to set both at once): `flow` (conditional flow matching, ODE-sampled in ~10 steps), `ddpm` (denoising diffusion), or `wgan` (single-pass WGAN-GP generator/critic).
+<details>
+<summary><h2 style="display:inline">Architecture</h2></summary>
 
-**Stage 1 — primary step.** Predicts the 9D post-step outcome (`giant/constants.py:LOCAL_TARGET_NAMES`) from the pre-step conditioning:
+**Stage 1 — primary step.** Predicts the 9D post-step outcome
+(`giant/constants.py:LOCAL_TARGET_NAMES`) from the pre-step conditioning:
 
 | Index | Variable | Encoding |
 |-------|----------|----------|
 | 0 | `step_length` [mm] | log |
-| 1–2 | `edep_logit`, `sec_logit` | ALR coords of the deposit/secondary/post-energy simplex |
+| 1–2 | `edep_logit`, `sec_logit` | ALR coordinates of the deposit / secondary / post-energy simplex |
-| 3–5 | `post_dir` in local frame | unit vector |
+| 3–5 | `post_dir` | unit vector, local frame (`pre_dir = ẑ`) |
-| 6–8 | `travel_dir` (`post_pos − pre_pos`) in local frame | unit vector |
+| 6–8 | `travel_dir` (`post_pos − pre_pos`) | unit vector, local frame |
 
-- Energy logits decode via softmax over `[edep_logit, sec_logit, 0]` × `pre_E`, so `edep + e_sec + post_E == pre_E` exactly — conservation is architectural, not learned.
+Energy logits decode via `softmax([edep_logit, sec_logit, 0]) × pre_E`, so
-- `post_dir`/`travel_dir` live in the frame where `pre_dir = ẑ`. `post_pos` isn't a target — it's reconstructed as `pre_pos + step_length * world_frame(travel_dir)`.
+`edep + e_sec + post_E == pre_E` holds exactly. `post_pos` is not itself a
+target — it's reconstructed as `pre_pos + step_length · world_frame(travel_dir)`,
+since duplicating that magnitude in a second target would let the two drift out
+of sync.
 
-**Stage 2 — secondaries.** Conditioned on the pre-step state and Stage 1's outcome, it generates the variable-length list of secondary particles. Two decoding strategies (`--stage2-decoder`):
+**Stage 2 — secondaries.** Conditioned on the pre-step state and Stage 1's
+outcome, it generates the variable-length secondary list, one token at a time in
+descending-energy order (`autoregressive`, default) or all `K_MAX` slots in one
+masked pass (`one_shot`). Autoregressive tokens condition on a running history —
+`markov` (previous token only) or `attention` (causal self-attention, KV-cached
+at inference). Either way, secondary energies stick-break the `e_sec` budget
+handed down from Stage 1, so the whole chain conserves energy. A secondary's
+species is represented `onehot` (categorical, top-N PDG codes + "other"),
+`physical` (continuous log-mass/charge), or `embedding` (nearest-neighbour
+lookup).
 
-- `autoregressive` — emits secondaries one at a time in descending-energy order, each token conditioned on a running history of prior tokens (`markov`: previous token only, or `attention`: causal self-attention, KV-cached at inference)
+**Conditioning (15D).** Pre-step position / energy / direction / layer, plus
-- `one_shot` — all `K_MAX` slots generated in a single forward pass, masked past the predicted `n_sec`
+particle mass/charge and material Z_eff/A_eff/density/X0/λ_int — encoded the
+same three ways as secondary species above, configured *independently* per axis
+(`conditioning.particle.type` / `conditioning.material.type`). `physical`
+computes rather than looks up, so it generalizes to species and materials
+outside the training menu; that's the default. `n_sec`/`e_sec` are always model
+outputs, never conditioning inputs.
 
-Either way, secondary energies stick-break the `e_sec` budget handed down from Stage 1, so the full chain conserves energy. A secondary's particle identity is represented as `onehot` (categorical, top-N PDG codes + "other"), `physical` (continuous log-mass/charge), or `embedding` (nearest-neighbour lookup).
+**Composable by design** — every stage assembles from small registries, so
+swapping one axis doesn't touch the others:
 
-**Conditioning.** Pre-step position/energy/direction/layer, plus particle mass/charge and material Z_eff/A_eff/density/X0/λ_int, encoded the same three ways as particle identity above. The particle and material axes are configured independently (`conditioning.particle.type` / `conditioning.material.type`; `--conditioning` sets both at once) and may mix — the `physical` representation generalizes to species/materials outside the training menu since it's computed rather than looked up. `n_sec`/`e_sec` are always model outputs, never conditioning inputs.
+| Registry | Choices |
+|---|---|
+| Objective | `flow` (matching, ~10-step ODE sample) · `ddpm` (denoising diffusion) · `wgan` (single-pass GAN) |
+| Trunk | `resmlp` · `none`, optionally MoE-routed (`RoutedTrunk`) |
+| Router | `energy` · `pdg` · `process` · `composed` · `none` — soft-mixed at train time, **top-1 dispatched at eval time**, which is the actual inference-speed win |
+| History (stage 2 AR) | `markov` · `attention` · `none` |
 
-**MoE routing** (`--router`): a pluggable `Router` (`energy`/`pdg`/`process`/`composed` axes) top-1-dispatches each row to one of several small expert trunks at eval time, instead of running one monolithic trunk. The CLI flags configure Stage 1's router; Stage 2 has its own `stage2_model.router` block, config-file only.
+</details>
 
-## Data
+<details>
+<summary><h2 style="display:inline">Configuration</h2></summary>
 
-- Input: parquet files produced by [miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from ROOT via `dwarf convert`. One row = one Geant4 step.
+Every default lives in one place: frozen dataclasses in `giant/config.py`,
-- **Conditioning (pre-step) columns:** `event_id`, `pdg`, `pre_x`/`pre_y`/`pre_z`, `pre_E`, `pre_dx`/`pre_dy`/`pre_dz` (direction), `material`, `layer_id`.
+composed into `GiantConfig` (`conditioning` / `stage1_model` / `stage2_model` /
-- **Primary outcome (post-step) columns:** `post_x`/`post_y`/`post_z`, `post_E`, `post_dx`/`post_dy`/`post_dz`, `step_length`, `edep` (energy deposited in this step), `e_sec` (total energy carried off by secondaries), `child_track_ids` (its length gives `n_sec`).
+`train`). `DEFAULT_CONFIG` is *generated* from `GiantConfig().to_dict()` rather
-- **Secondary columns**, one variable-length list per step: `sec_pdg_list`, `sec_E_list`, `sec_dx_list`/`sec_dy_list`/`sec_dz_list` — padded/truncated to `K_MAX` (15) slots on load, ordered by descending energy.
+than hand-maintained, so the dataclasses can't drift from what actually gets
-- **Optional:** `process` — the physics process that produced the step (e.g. `compt`, `phot`, `eBrem`); a post-step label used only as classifier supervision (`ProcessRouter`), never as conditioning.
+merged. TOML config keys are validated against that shape — an unknown key is
-- Train/val split is by `event_id` (`--seed`-controlled), not row shuffle, so correlated steps from the same shower never leak across the split.
+rejected with a did-you-mean suggestion. Precedence: CLI flag > `--config` file
-- Loading a directory or `.manifest` of multiple parquet files (each one Geant4 job, `event_id` restarting from 0) offsets each file's `event_id`s by a fixed per-file stride so ids stay globally unique across files.
+> default.
 
-## Project structure
+```toml
+# config.toml — resolved shape of the four blocks
+[conditioning]
+particle.type = "physical"
+material.type = "physical"
 
-```
+[stage1_model]
-giant/
+generator = "flow"
-├── giant/
+
-│   ├── data/
+[stage2_model]
-│   │   ├── loader.py       # parquet → numpy arrays (incl. streaming/chunked reads)
+generator = "wgan"
-│   │   ├── transforms.py   # log transforms, local-frame rotation, energy simplex, secondary encode/decode
+decoder   = "autoregressive"
-│   │   ├── dataset.py      # StepsDataset / StreamingStepsDataset (PyTorch)
+
-│   │   └── setup_cache.py  # sidecar cache for the pre-epoch setup scan (vocab/split/normalizers)
+[train]
-│   ├── model/
+epochs     = 100
-│   │   ├── models.py       # Stage1Model, Stage2OneShot, Stage2Autoregressive, CriticModel
+batch_size = 4096
-│   │   ├── builders.py     # build_models / build_critics — config dict → assembled stage models
+lr         = 3e-4
-│   │   ├── encoders.py     # ConditionEncoder (physical / embedding / onehot, per axis)
-│   │   ├── layers.py       # ResBlock/AdaLNResBlock registry, SinusoidalEmbedding, MLP heads
-│   │   ├── trunks.py       # trunk registry (resmlp, none) + RoutedTrunk (MoE expert bodies)
-│   │   ├── routers.py      # Router registry: energy / pdg / process / composed / none
-│   │   ├── history.py      # stage-2 AR history encoders: markov / attention (KV-cached) / none
-│   │   ├── objectives.py   # flow / ddpm / wgan objective registry
-│   │   ├── schedule.py     # CosineSchedule (DDPM) and flow matching utilities
-│   │   ├── wgan.py         # WGAN-GP gradient penalty / critic / generator losses
-│   │   ├── summary.py      # build-only introspection behind `giant model summary`
-│   │   ├── _legacy.py      # v0.2 checkpoint model_config/state-dict migration
-│   │   └── network.py      # re-export shim over all of the above
-│   ├── constants.py        # output/conditioning dims, K_MAX, secondary slot layout, schema keys
-│   ├── cond_layout.py      # single source of truth for the cond_cont/cond_cat column layout
-│   ├── particles.py        # PDG → (mass, charge) decode, incl. nuclear/ion codes; onehot/embedding secondary-identity decode
-│   ├── materials.py        # material name → (Z_eff, A_eff, density, X0, λ_int)
-│   ├── config.py           # default hyperparameters, TOML config merging, device autodetect
-│   ├── pipeline.py         # builds datasets/normalizers and kicks off a training run (with setup-stage caching)
-│   ├── training/           # two-stage training: loop, per-stage trainers, metrics, checkpointing
-│   │   ├── loop.py         #   epoch loop, graceful shutdown, best-checkpoint selection
-│   │   ├── trainers.py     #   StageSpec + flow/ddpm and WGAN-GP per-stage trainers
-│   │   ├── stage2_inputs.py#   ground-truth stage-2 targets + autoregressive/teacher-forcing inputs
-│   │   ├── metrics.py      #   MetricsCollector: metrics.csv columns, W&B logging, progress/summary
-│   │   ├── amp.py          #   bf16 autocast (`train.precision`)
-│   │   ├── plots.py        #   training-progress plots (`giant analyze metrics`)
-│   │   └── checkpoint.py   #   checkpoint assembly/restore (format unchanged since v0.2)
-│   ├── sample.py           # DDPM / DDIM / flow matching / WGAN samplers + secondary sampling
-│   ├── checkpoint_io.py    # checkpoint → ready-to-run models/normalizers (predict + rollout)
-│   ├── geometry.py         # GeometryOracle: position → (material, layer_id, escaped) for rollout
-│   ├── rollout.py          # autoregressive shower rollout driver
-│   ├── validate.py         # step-level marginal + KL-divergence validation
-│   ├── _migration.py       # shared v0.2 → v0.3 facts used by both migration surfaces
-│   ├── analysis/           # rollout-vs-reference analysis pipeline (see `giant analyze` below)
-│   │   ├── sources.py      #   canonical LazyFrames + secondary view
-│   │   ├── variables.py    #   per-step value expressions shared by range sizing and the catalog
-│   │   ├── reduce.py       #   streaming reduction primitives (hist1d, per-event scalars, profiles, ...)
-│   │   ├── grouping.py     #   fixed bin edges + energy/pdg/material group sets
-│   │   ├── context.py      #   resolves grouping into `shared.json` once per run
-│   │   ├── reduced.py      #   Partial/Reduced — the compact JSON a compute job emits
-│   │   ├── catalog.py      #   declarative PlotSpec registry (`giant analyze list`)
-│   │   ├── router_gating.py / type_embedding_distance.py   # checkpoint-bound diagnostics
-│   │   ├── runtime_estimate.py  #   per-(plot, chunk) walltime estimates for submit
-│   │   ├── condor.py       #   prep / compute-one / merge / submit-description plumbing
-│   │   └── render.py       #   PDFs + HTML gallery (only module importing plotstyle/LaTeX)
-│   └── cli.py              # `giant train` / `new-run` / `model summary` / `predict` / `rollout` / `analyze`
-├── giant/tools/                        # dataset/tooling logic, unified under the `dwarf` CLI (`dwarf --help`)
-│   ├── dwarf.py                        # Typer app: convert, migrate, bump-gen, bump-schema, status,
-│   │                                    #   update-manifest, create-manifest, make-root,
-│   │                                    #   build-geometry-oracle, warm-cache, hparam-scan
-│   ├── steps_to_parquet.py             # ROOT → parquet conversion (uproot/awkward/polars) — `dwarf convert`
-│   ├── steps_to_parquet_parallel.py    # fan out conversion over several ROOT files — `dwarf convert --jobs N`
-│   ├── migrate_geant_steps.py          # one-time move into the raw/processed/pools/derived layout — `dwarf migrate`
-│   ├── bump_dataset_version.py         # cut a new raw gen or parquet schema, with a logged reason —
-│   │                                    #   `dwarf bump-gen` / `bump-schema` / `status` / `update-manifest` / `create-manifest`
-│   ├── create_root_files.py            # generate new ROOT shards via a minicalosim executable — `dwarf make-root`
-│   ├── geometry_oracle.py              # fit a position → (material, layer_id) oracle — `dwarf build-geometry-oracle`
-│   ├── warm_setup_cache.py             # precompute `giant train`'s setup-stage sidecar — `dwarf warm-cache`
-│   ├── hparam_scan.py                  # hyperparameter grid scan over `giant train` runs — `dwarf hparam-scan`
-│   └── profile_analysis_costs.py       # profiling helper for the `giant analyze` reduction pipeline
-└── tests/
 ```
 
-## Setup
+Some knobs only exist in the config file, with no CLI flag:
+`stage2_model.autoregressive.teacher_forcing`/`.history`,
+`stage2_model.particle_type.target`/`.class_weighting`,
+`stage2_model.n_sec.mode`/`.owner`, `conditioning.share_stages`,
+`stage2_model.router.*`, and the finer `router` knobs (`lambda_balance`,
+`gumbel`, `learn_width`, …).
+
+`configs/` holds kept reference configs — `baseline.toml` is the fixed
+comparison point every experimental variant (routed trunk, WGAN, attention
+history, embedding conditioning) is a single edit away from. v0.2 flat-schema
+configs and checkpoints load and auto-migrate.
+
+</details>
+
+<details>
+<summary><h2 style="display:inline">Data</h2></summary>
+
+Input is parquet — one row per Geant4 step — from
+[miniCaloSim](https://gitlab.etp.kit.edu/lbogner/minicalosim), or converted from
+ROOT via `dwarf convert`.
+
+| Group | Columns |
+|---|---|
+| **Conditioning (pre-step)** | `event_id`, `pdg`, `pre_x`/`pre_y`/`pre_z`, `pre_E`, `pre_dx`/`pre_dy`/`pre_dz`, `material`, `layer_id` |
+| **Primary outcome (post-step)** | `post_x`/`post_y`/`post_z`, `post_E`, `post_dx`/`post_dy`/`post_dz`, `step_length`, `edep`, `e_sec`, `child_track_ids` (length → `n_sec`) |
+| **Secondaries** (variable-length lists) | `sec_pdg_list`, `sec_E_list`, `sec_dx_list`/`sec_dy_list`/`sec_dz_list` — padded/truncated to `K_MAX = 15` slots, descending energy |
+| **Optional** | `process` — physics-process label, classifier supervision only (`ProcessRouter`), never conditioning |
+
+Train/val split is by `event_id` (`--seed`-controlled), not row shuffle, so a
+shower's correlated steps never straddle the split. Loading a directory or
+`.manifest` of several parquet files offsets each file's `event_id`s by a
+per-file stride so ids stay globally unique. The pre-epoch setup scan (vocab
+maps, event split, normalizer stats) persists to a sidecar cache
+(`--cache-setup`/`--rebuild-setup-cache`), precomputable ahead of time via
+`dwarf warm-cache`.
+
+</details>
+
+<details>
+<summary><h2 style="display:inline">CLI reference</h2></summary>
+
+**`giant`** — train, run, and analyze the surrogate:
+
+| Command | Does |
+|---|---|
+| `new-run` | scaffold a `config.toml` + run directory from flags |
+| `train DATA` | train the two-stage model |
+| `model summary` | build-only parameter counts, without training |
+| `predict DATA --checkpoint …` | per-step predictions from a checkpoint |
+| `rollout DATA --checkpoint … --geometry …` | full autoregressive shower rollout |
+| `analyze prep/submit` | build a run dir; `submit` also queues HTCondor compute jobs |
+| `analyze compute-one` / `merge-one` | one plot × chunk reduction / merge (what a condor job runs) |
+| `analyze render <run_dir> --gallery` | merge chunks → styled PDFs + HTML gallery (local, needs LaTeX) |
+| `analyze metrics <train_run_dir>` | training-progress plots from `metrics.csv` |
+| `analyze list` | every catalog plot id |
+
+**`dwarf`** — dataset/tooling CLI:
+
+| Command | Does |
+|---|---|
+| `convert` | ROOT Steps tree → parquet (`--jobs N` fans out) |
+| `migrate` | one-time move into the raw/processed/pools/derived layout |
+| `bump-gen` / `bump-schema` / `status` | dataset versioning |
+| `update-manifest` / `create-manifest` | point/build a manifest of parquet files |
+| `make-root` | generate new ROOT shards via a minicalosim executable |
+| `build-geometry-oracle` | fit position → (material, layer_id) for rollout |
+| `warm-cache` | precompute `giant train`'s setup-stage sidecar |
+| `hparam-scan` | grid-scan dropout × n_blocks × hidden_dim |
+
+Worth knowing on `giant train` (full surface behind `--help`):
+`--mode {flow,ddpm,wgan}` / `--stage1-generator` / `--stage2-generator`,
+`--stage2-decoder {autoregressive,one_shot}`, `--conditioning
+{physical,embedding,onehot}`, `--router` / `--router-type` / `--n-experts` /
+`--router-axis`, `--stage{1,2}-init-from` + `--stage{1,2}-freeze` (retrain one
+stage against a fixed other one), `--precision {fp32,bf16}`, `--wandb`.
+
+</details>
+
+<details>
+<summary><h2 style="display:inline">Rollout &amp; analysis</h2></summary>
+
+`giant rollout` seeds showers from each event's highest-energy entry step, then
+autoregressively steps the model to completion — advancing all active tracks
+breadth-first, batched — pushing secondaries as new tracks and looking up
+`material`/`layer_id` from the geometry oracle each step. Tracks terminate on
+one of six reasons (energy cutoff, max steps, detector escape, natural end,
+unknown pdg, max tracks); every reason but escape deposits the remaining energy
+locally, so showers conserve energy by construction — only `escaped` counts as
+leakage.
+
+`giant analyze` compares one or more rollouts against a single held-out
+reference: `prep` resolves shared bin edges/groups once, `submit`/`compute-one`
+run each (plot, `event_id`-disjoint chunk) pair as a polars/numpy-only HTCondor
+job, `render` merges the chunks and produces the styled PDFs + HTML gallery
+locally (the only step that needs LaTeX). Each rollout gets its own colored
+series against one shared reference line. `giant analyze metrics` is a separate
+entry point — training-progress plots straight from a run's `metrics.csv`.
+
+</details>
+
+<details>
+<summary><h2 style="display:inline">Install</h2></summary>
+
+| Extra | Adds | For |
+|---|---|---|
+| `cpu` **or** `cuda` | torch 2.3.x | required — mutually exclusive, pick one |
+| `geometry` | scikit-learn | `dwarf build-geometry-oracle`, rollout |
+| `analysis` | matplotlib, plotstyle | `giant analyze render` |
+| `convert` | uproot, awkward | `dwarf convert` |
+| `wandb` | wandb | `giant train --wandb` |
+| `dev` | pytest, ruff, ty, + all of the above | development |
 
 ```bash
-uv sync --extra cpu                   # CPU-only torch (use --extra cuda for CUDA 11.8 instead)
+uv sync --extra cpu --extra dev     # everything needed to develop
-uv sync --extra cpu --extra dev       # add dev tools (pytest, ruff, ty)
-uv sync --extra cpu --extra geometry  # add scikit-learn, for `dwarf build-geometry-oracle` / rollout
-uv sync --extra cpu --extra analysis  # matplotlib/polars/plotstyle, for `giant analyze render`
-uv sync --extra cpu --extra convert   # uproot/awkward/polars, for `dwarf convert`
-uv sync --extra cpu --extra wandb     # W&B logging (`giant train --wandb`)
 ```
 
-The `dev` extra pulls in `convert`, `analysis`, `geometry` and `wandb` as well.
+Plain `uv sync` with no extra installs **no torch at all** — always include
+`--extra cpu` or `--extra cuda`.
 
-`cpu` and `cuda` are mutually exclusive — pick one to select the torch build (pinned to 2.3.x). Plain `uv sync` installs no torch at all. See `CLAUDE.md` for details.
+</details>
 
-## Training, prediction, rollout
+<details>
+<summary><h2 style="display:inline">Development</h2></summary>
 
 ```bash
-giant new-run --hidden-dim 512 --lr 3e-4 --comment "..."   # scaffold a config.toml + run dir
+uv run pytest            # 964 tests
-giant train path/to/steps.parquet                          # train (flow stage 1 + wgan stage 2, default)
-giant predict path/to/steps.parquet --checkpoint checkpoints/.../best.pt
-
-dwarf build-geometry-oracle path/to/steps.parquet --out oracle.pkl   # position → material/layer_id
-giant rollout path/to/steps.parquet --checkpoint checkpoints/.../best.pt --geometry oracle.pkl
-```
-
-Useful flags on `giant train`:
-
-- `--mode {flow,ddpm,wgan}` sets both stages' objective at once; `--stage1-generator`/`--stage2-generator` override per stage
-- `--stage2-decoder {autoregressive,one_shot}` — Stage 2 decoding strategy (see Architecture)
-- `--conditioning {physical,embedding,onehot}` — conditioning representation
-- `--router` / `--router-type` / `--n-experts` / `--router-axis` — MoE routing
-- `--stage2-stage1-context {truth,sampled}` — feed Stage 2 the ground-truth or the model's own sampled Stage-1 outcome (annealable via `stage2_model.ctx_p_start`/`ctx_p_end`)
-- `--precision {fp32,bf16}` — bf16 autocast in the training loop
-- `--wandb` — log per-epoch metrics to Weights & Biases (needs `uv sync --extra wandb`); metric names are `<stage>/<split>/<metric>` plus an unprefixed run-level tail, all derived from `giant/training/trainers.py` `MetricSpec`s
-- `--no-cache-setup` / `--rebuild-setup-cache` — control the setup-stage sidecar cache (vocab maps, event split, normalizer stats); `dwarf warm-cache` precomputes it
-- `--stage1-init-from`/`--stage2-init-from` (checkpoint `.pt`) + `--stage1-freeze`/`--stage2-freeze` — load a stage's weights from another checkpoint and never update them, so the other stage can be retrained alone against a fixed, known-good one while still producing a complete, rollout-capable checkpoint
-
-Config-file-only knobs (no CLI flag — use `--config config.toml`): `stage2_model.autoregressive.teacher_forcing`/`.history`, `stage2_model.particle_type.target`/`.class_weighting`, `stage2_model.n_sec.mode`/`.owner`, `conditioning.share_stages`, `stage*_model.trunk.*` and the finer `router` knobs (`lambda_balance`, `gumbel`, `learn_width`, …). `configs/` holds kept reference configs. v0.2 flat-schema configs and checkpoints load fine (auto-migrated).
-
-`giant rollout` seeds showers from each event's highest-energy entry step, then autoregressively steps the model to completion, pushing secondaries as new tracks and looking up `material`/`layer_id` from the geometry oracle each step. Tracks terminate on energy cutoff, max steps, detector escape, or natural end; energy is deposited locally on every stop except escape, so showers conserve energy by construction.
-
-## Validation and analysis
-
-- `giant.validate.validate_marginals` — step-level marginal + KL-divergence checks during training (`--validate-every`)
-- `giant analyze` — deeper rollout-vs-reference diagnostics (marginals by energy/pdg/material, per-event totals, shower profiles, species share, leakage, secondaries):
-
-```bash
-giant analyze submit rollout.yaml --accounting-group cms   # prep + one HTCondor job per plot × chunk (compute only)
-giant analyze submit a.yaml b.yaml --accounting-group cms --label flow --label wgan  # N rollouts vs one shared reference
-giant analyze render <run_dir> --gallery                   # local: merge chunks, then styled PDFs + HTML gallery (needs LaTeX)
-
-giant analyze list                                         # every catalog plot id
-giant analyze prep rollout.yaml --chunks 8                 # just the run directory, no submission
-giant analyze compute-one --id marginal_edep --run-dir <run_dir> --chunk 0   # what a condor job runs
-giant analyze merge-one --id marginal_edep --run-dir <run_dir>               # merge one plot's chunks (debugging)
-```
-
-`<run_dir>` defaults to `<cwd>/analysis_runs/analysis_<id>` (`--run-dir` overrides it; `prep`/`submit` print it). Multiple rollout YAMLs must all name the same reference (`dataset`) file; each renders as its own colored series against one reference line/panel. Compute jobs are polars/numpy only; only `render` needs LaTeX, so it always runs locally.
-
-Separately, `giant analyze metrics <train_run_dir>` renders training-progress plots (loss/lr/accuracy/grad-norm/router/wgan/throughput) straight from a training run's `metrics.csv`.
-
-## Development
-
-```bash
-uv run pytest            # run tests
 uv run ruff check .      # lint
 uv run ruff format .     # format
 uv run ty check .        # type check
 ```
+
+Gitea Actions (`.gitea/workflows/ci.yml`) runs lint + format-check + type-check
++ tests on every push and PR; merges to `master` auto-bump the patch version
+and regenerate `CHANGELOG.md` — don't hand-edit either.
+
+</details>
+
+## License
+
+Not yet decided — treat this repository as all-rights-reserved until a
+`LICENSE` file is added.

configs/baseline.toml

@@ -84,15 +84,30 @@ dropout      = 0.0
 # secondary-species failure. Flow (not the schema default wgan) so the
 # baseline varies only the decoder relative to the best v0.2 result.
 #
-# COST, measured (RTX 4070, bs 4096, 10 ODE steps), not estimated:
+# COST, measured (RTX 4070, bs 4096, 10 ODE steps), not estimated — but see
-# sample.sample_secondaries_ar loops `for k in range(k_max)` unconditionally
+# the row-compaction note below, which changes the INFERENCE side of this:
-# — all 15 slots regardless of predicted n_sec — so a flow AR token costs
-# k_max * steps = 150 stage-2 calls per physics step. That makes this block
-# the dominant cost on both sides:
 #   training   flow AR 29.5k samp/s  vs flow one-shot 190.7k samp/s  (6.5x)
 #   inference  flow AR  8.5k step/s  vs flow one-shot  68.7k step/s  (8.1x)
 # Accepted deliberately: one-shot is the configuration whose secondary
 # species distribution failed, and that failure is what v0.3 exists to fix.
+#
+# Row compaction (landed after the above measurement): at inference,
+# sample.sample_secondaries_ar used to loop `for k in range(k_max)`
+# unconditionally — all 15 slots regardless of predicted n_sec — so a flow
+# AR token cost k_max * steps = 150 stage-2 calls per physics step. It now
+# drops a row from the batch the moment its own secondary count is
+# exhausted, so the real inference cost is ~n_sec * steps stage-2 calls
+# (this checkpoint's own rollout measured 0.382 secondaries/step — see
+# giant-baseline-flow-ar-rollout-validation.md), not k_max * steps. A CPU
+# micro-benchmark at that multiplicity (giant/model/history.py's
+# hidden_dim=512/6-block shape, k_max=15, batch 512) measured 17.6-22.9x
+# fewer wall-clock seconds for the AR loop alone (markov/attention history
+# respectively) — directional only (CPU, synthetic n_sec distribution, not
+# an end-to-end rollout); the 8.1x inference ratio above is now stale and
+# should be re-measured on GPU via a real rollout + `eval_cost_per_step`
+# once one is run against this checkpoint. Training cost (the 6.5x/29.5k
+# figures) is untouched by this: teacher_forcing = "always" here never
+# calls the AR sampler at train time (see [stage2_model.autoregressive]).
 decoder      = "autoregressive"
 generator    = "flow"
 hidden_dim   = 512

giant/model/history.py

@@ -4,6 +4,7 @@ for the `HISTORY_REGISTRY`/`build_history` factory, which mirrors
 `giant.model.routers`'s `Router`/`ROUTER_REGISTRY` pattern (gitea #35)."""
 
 import inspect
+from typing import cast
 
 import torch
 import torch.nn as nn
@@ -34,6 +35,17 @@ class HistoryEncoder(nn.Module):
     def step(self, feat: torch.Tensor, has_prev: torch.Tensor, cache: object) -> tuple[torch.Tensor, object]:
         return self.forward(feat, has_prev), cache
 
+    def select_cache(self, cache: object, idx: torch.Tensor) -> object:
+        """Row-compacts an inference cache (`init_cache`/`step`'s state) down
+        to `idx` — used by `giant.sample.sample_secondaries_ar`'s row
+        compaction to keep a shrinking active-row set's cache aligned as rows
+        finish generating. Default here matches `init_cache`/`step`'s O(1)
+        default: `cache` is always `None`, so there's nothing to index —
+        correct for any encoder whose per-step state doesn't carry a batch
+        dimension (`MarkovHistory` has no cache at all; its running state is
+        `prev_repr`/`remaining`, compacted directly by the caller)."""
+        return cache
+
 
 HISTORY_REGISTRY: dict[str, type[HistoryEncoder]] = {}
 
@@ -215,3 +227,13 @@ class AttentionHistory(HistoryEncoder):
             x, kv_new = block.step(x, kv)
             new_cache.append(kv_new)
         return x, new_cache
+
+    def select_cache(self, cache: object, idx: torch.Tensor) -> list[torch.Tensor | None]:
+        """Row-compacts every block's `(B, T, dim)` KV cache down to `idx`
+        along its batch dimension — see `HistoryEncoder.select_cache`. `idx`
+        may be a long index tensor or a boolean mask (`giant.sample`'s AR
+        loop uses both). `None` entries (a block that has never seen a
+        `step` call yet) stay `None`."""
+        assert isinstance(cache, list)
+        cache_t = cast("list[torch.Tensor | None]", cache)
+        return [None if kv is None else kv[idx] for kv in cache_t]

giant/model/models.py

@@ -614,6 +614,14 @@ class Stage2Autoregressive(StageModel):
         slot — see `AttentionHistory.step`'s docstring."""
         return self.history_encoder.step(token_feat, has_prev, cache)
 
+    def select_history_cache(self, cache, idx: torch.Tensor):
+        """Row-compacts `cache` (from `init_history_cache`/`history_step`)
+        down to `idx` — see `HistoryEncoder.select_cache`. Used by
+        `giant.sample.sample_secondaries_ar`'s active-row compaction to keep
+        the cache aligned with a shrinking batch as rows finish generating
+        across AR slots."""
+        return self.history_encoder.select_cache(cache, idx)
+
     def forward(
         self,
         x_t: torch.Tensor,

giant/sample.py

@@ -231,6 +231,7 @@ def sample_secondaries_ar(
     stage1_out: torch.Tensor,
     n_sec_pred: torch.Tensor | None,
     steps: int = 10,
+    full_length: bool = False,
 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
     """`Stage2Autoregressive` inference loop: one token at a time, in
     descending-energy slot order, up to `k_max` sequential calls. Unlike
@@ -242,19 +243,18 @@ def sample_secondaries_ar(
     expressiveness.
 
     A `{flow,ddpm}` token costs `steps` ODE substeps; `wgan` costs one pass —
-    the "K sequential forwards" cost applies per-token here, not
+    the "K sequential forwards" cost applies per-token here, not once, so a
-    once, so a flow/ddpm AR run costs ~`k_max * steps` model calls per
+    flow/ddpm AR run costs ~`n_sec * steps` model calls per physics step
-    physics step (or ~`n_sec * steps` under `n_sec_pred=None` below, once
+    (measured on `configs/baseline.toml`: 0.382 secondaries/step at rollout
-    every row in the batch has stopped).
+    time), not `k_max * steps` — see the row-compaction paragraph below.
 
-    `n_sec_pred`, if given, fixes each row's secondary count up front (as
+    `n_sec_pred`, if given (as resolved by `resolve_n_sec` — `n_sec.mode` in
-    resolved by `resolve_n_sec` — `n_sec.mode` in `("head", "truth")`, or a
+    `("head", "truth")`, or a stop-token decoder driven by
-    stop-token decoder driven by `_assemble_stage2_ar_inputs_scheduled`'s
+    `_assemble_stage2_ar_inputs_scheduled`'s ground-truth `n_sec`, which must
-    ground-truth `n_sec`, which must run the *full* `k_max`-length free-
+    run the *full* `k_max`-length free-running self-sample regardless of the
-    running self-sample regardless of the decoder's own stop head — the
+    decoder's own stop head — the scheduled-sampling training contract does
-    scheduled-sampling training contract does not truncate). This always
+    not truncate, see `full_length` below) fixes each row's secondary count
-    runs the full `k_max`-iteration loop, masking by the given count at the
+    up front.
-    end exactly as before.
 
     `n_sec_pred=None` is only valid when `sec_decoder.stop_head` is set
     (`n_sec.mode = "stop_token"`): before generating each slot's token, that
@@ -263,11 +263,35 @@ def sample_secondaries_ar(
     why this needs no extra state) decides whether generation should have
     already stopped, per `sec_decoder.n_sec_sampling` ("greedy": threshold at
     0; "sample": a Bernoulli draw at `sigmoid(logit)`). A row's own
-    `n_sec_pred` is the first slot index where this fires; once every row in
+    `n_sec_pred` is the first slot index where this fires. A row that never
-    the batch has fired, the loop breaks before spending a model call on the
+    fires within `k_max` is capped there (`K_MAX` stays a safety cap, not a
-    next slot's token — the average-case cost win the docstring above
+    modeling ceiling).
-    describes. A row that never fires within `k_max` is capped there
+
-    (`K_MAX` stays a safety cap, not a modeling ceiling).
+    **Row compaction.** A row that has already produced its `n_sec_pred`
+    tokens (or, under `stop_token`, has already fired its stop logit) has
+    nothing left to contribute — every later slot of that row is masked out
+    of `sec_valid` on return, and downstream consumers (`giant/rollout.py`,
+    `giant/cli.py`) never read it. So unless `full_length=True`, this
+    function drops such rows from the active set entirely instead of running
+    the model on them: `active_idx` starts at every row with `n_sec_pred > 0`
+    (or, under `stop_token`, every row — the first stop decision can fire at
+    slot 0) and only shrinks as rows finish, so slot `k`'s model calls cost
+    `O(active rows)` not `O(B)`. Slots a row never reaches keep their `0.0`
+    zero-init in `sec_cont`/`sec_type` — masked by `sec_valid`, identical to
+    what a full, uncompacted run would have written there before masking.
+    `AttentionHistory`'s KV cache is kept aligned to the shrinking active set
+    via `Stage2Autoregressive.select_history_cache`
+    (`giant.model.history.HistoryEncoder.select_cache`) every time the set
+    shrinks; `MarkovHistory`'s O(1) state (`prev_repr`/`remaining`, carried
+    directly rather than through a cache) is compacted the same way.
+
+    `full_length=True` disables all of the above: every row runs the full
+    `k_max`-iteration loop regardless of `n_sec_pred`/stop decisions, exactly
+    reproducing the pre-compaction behaviour. Required by
+    `_assemble_stage2_ar_inputs_scheduled`'s scheduled-sampling self-sample,
+    whose training contract needs a real prediction at every slot up to
+    `k_max` (mixed per-slot against ground truth) even past a row's own
+    `n_sec` — see that function's docstring.
 
     Under `history="attention"` the history encoding is computed once per
     slot via `Stage2Autoregressive.history_step` (a KV-cache append)
@@ -312,11 +336,6 @@ def sample_secondaries_ar(
     sec_cont = torch.zeros(B, k_max, CONT_SLOT_DIM, device=device)
     sec_type = torch.zeros(B, k_max, type_dim, device=device)
 
-    # Running per-token state, threaded from one slot to the next.
-    prev_repr = torch.zeros(B, CONT_SLOT_DIM + type_dim, device=device)
-    remaining = torch.ones(B, device=device)
-    history_cache = sec_decoder.init_history_cache()
-
     use_stop_token = n_sec_pred is None
     if use_stop_token:
         assert getattr(sec_decoder, "stop_head", None) is not None, (
@@ -324,21 +343,46 @@ def sample_secondaries_ar(
             "with no stop_head — only valid under stage2_model.n_sec.mode = "
             "'stop_token'"
         )
-        finished = torch.zeros(B, dtype=torch.bool, device=device)
         derived_n_sec = torch.full((B,), k_max, dtype=torch.long, device=device)
+        # Tracks which rows have already recorded a stop, globally by
+        # original batch index — needed even under compaction's own
+        # never-revisit guarantee, because `full_length=True` keeps every
+        # row in `active_idx` for the whole loop, so a row whose stop logit
+        # fires once but flips back below threshold at a later slot (a real
+        # possibility for an untrained/lightly-trained stop_head) must not
+        # have `derived_n_sec` overwritten by that later, spurious re-fire.
+        finished = torch.zeros(B, dtype=torch.bool, device=device)
+
+    # `active_idx`: rows still contributing tokens, indexed into the
+    # original batch. Only ever shrinks (never full_length) or stays fixed
+    # at arange(B) (full_length) — see the row-compaction docstring section.
+    active_idx = torch.arange(B, device=device)
+    if not full_length and not use_stop_token:
+        active_idx = active_idx[n_sec_pred > 0]
+
+    # Running per-token state, already compacted to `active_idx`.
+    prev_repr = torch.zeros(active_idx.numel(), CONT_SLOT_DIM + type_dim, device=device)
+    remaining = torch.ones(active_idx.numel(), device=device)
+    history_cache = sec_decoder.init_history_cache()
 
     for k in range(k_max):
-        has_prev = torch.full((B, 1), k >= 1, dtype=torch.bool, device=device)
+        if active_idx.numel() == 0:
-        history_feat = prev_repr.unsqueeze(1)  # (B, 1, CONT_SLOT_DIM + type_dim)
+            break
-        remaining_frac = remaining.unsqueeze(1)  # (B, 1)
+        Bc = active_idx.numel()
-        slot_idx = torch.full((B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32)
+        cc = cond_cont.index_select(0, active_idx)
+        ck = cond_cat.index_select(0, active_idx)
+        s1 = stage1_out.index_select(0, active_idx)
+        has_prev = torch.full((Bc, 1), k >= 1, dtype=torch.bool, device=device)
+        history_feat = prev_repr.unsqueeze(1)  # (Bc, 1, CONT_SLOT_DIM + type_dim)
+        remaining_frac = remaining.unsqueeze(1)  # (Bc, 1)
+        slot_idx = torch.full((Bc, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32)
         hist, history_cache = sec_decoder.history_step(history_feat, has_prev, history_cache)
 
         if use_stop_token:
             stop_logit = sec_decoder.predict_stop(
-                cond_cont,
+                cc,
-                cond_cat,
+                ck,
-                stage1_out,
+                s1,
                 history_feat,
                 has_prev,
                 remaining_frac,
@@ -346,21 +390,31 @@ def sample_secondaries_ar(
                 hist=hist,
             ).squeeze(1)
             if sec_decoder.n_sec_sampling == "sample":
-                stop_now = torch.rand(B, device=device) < torch.sigmoid(stop_logit)
+                stop_now = torch.rand(Bc, device=device) < torch.sigmoid(stop_logit)
             else:
                 stop_now = stop_logit >= 0.0
-            derived_n_sec[stop_now & ~finished] = k
+            newly_stopped = stop_now & ~finished.index_select(0, active_idx)
-            finished = finished | stop_now
+            derived_n_sec[active_idx[newly_stopped]] = k
-            if finished.all():
+            finished[active_idx[stop_now]] = True
-                break
+            if not full_length:
+                keep = ~stop_now
+                active_idx = active_idx[keep]
+                cc, ck, s1 = cc[keep], ck[keep], s1[keep]
+                has_prev, remaining_frac, slot_idx = has_prev[keep], remaining_frac[keep], slot_idx[keep]
+                history_feat, hist = history_feat[keep], hist[keep]
+                history_cache = sec_decoder.select_history_cache(history_cache, keep)
+                prev_repr, remaining = prev_repr[keep], remaining[keep]
+                if active_idx.numel() == 0:
+                    break
+                Bc = active_idx.numel()
 
         if objective.is_adversarial:
-            z = torch.randn(B, 1, sec_decoder.noise_dim, device=device)
+            z = torch.randn(Bc, 1, sec_decoder.noise_dim, device=device)
             token = sec_decoder(
                 z,
-                cond_cont,
+                cc,
-                cond_cat,
+                ck,
-                stage1_out,
+                s1,
                 history_feat,
                 has_prev,
                 remaining_frac,
@@ -368,15 +422,15 @@ def sample_secondaries_ar(
                 hist=hist,
             )
         else:
-            x = torch.randn(B, 1, token_dim, device=device)
+            x = torch.randn(Bc, 1, token_dim, device=device)
             dt = 1.0 / steps
             for i in range(steps):
-                t = torch.full((B, 1), i * dt, device=device)
+                t = torch.full((Bc, 1), i * dt, device=device)
                 v = sec_decoder(
                     x,
-                    cond_cont,
+                    cc,
-                    cond_cat,
+                    ck,
-                    stage1_out,
+                    s1,
                     history_feat,
                     has_prev,
                     remaining_frac,
@@ -387,15 +441,15 @@ def sample_secondaries_ar(
                 x = x + v * dt
             token = x
 
-        token = token.squeeze(1)  # (B, token_dim)
+        token = token.squeeze(1)  # (Bc, token_dim)
         cont_k = token[:, :CONT_SLOT_DIM]
         if type_folded:
             type_k = token[:, CONT_SLOT_DIM:]
         else:
             type_k = sec_decoder.predict_type(
-                cond_cont,
+                cc,
-                cond_cat,
+                ck,
-                stage1_out,
+                s1,
                 history_feat,
                 has_prev,
                 remaining_frac,
@@ -403,8 +457,8 @@ def sample_secondaries_ar(
                 hist=hist,
             ).squeeze(1)
 
-        sec_cont[:, k] = cont_k
+        sec_cont[active_idx, k] = cont_k
-        sec_type[:, k] = type_k
+        sec_type[active_idx, k] = type_k
 
         if target == "onehot":
             type_for_history = F.one_hot(type_k.argmax(dim=-1), num_classes=type_dim).float()
@@ -415,6 +469,12 @@ def sample_secondaries_ar(
         prev_repr = torch.cat([stick_fraction.unsqueeze(-1), cont_k[:, 1:4], type_for_history], dim=-1)
         remaining = torch.clamp(remaining * (1.0 - stick_fraction), min=0.0)
 
+        if not full_length and not use_stop_token:
+            keep2 = n_sec_pred.index_select(0, active_idx) > (k + 1)
+            active_idx = active_idx[keep2]
+            prev_repr, remaining = prev_repr[keep2], remaining[keep2]
+            history_cache = sec_decoder.select_history_cache(history_cache, keep2)
+
     resolved_n_sec = derived_n_sec if use_stop_token else n_sec_pred
     sec_valid = torch.arange(k_max, device=device).unsqueeze(0) < resolved_n_sec.unsqueeze(1)
     return sec_cont, sec_type, sec_valid

giant/training/stage2_inputs.py

@@ -291,9 +291,13 @@ def _assemble_stage2_ar_inputs_scheduled(
     skips self-sampling entirely), so callers can call this unconditionally.
 
     The free-running estimate is a REAL autoregressive self-sample —
-    `giant.sample.sample_secondaries_ar` under `torch.no_grad()` — not a
+    `giant.sample.sample_secondaries_ar` under `torch.no_grad()`, called here
-    cheap one-step proxy, so building it costs the same `k_max` (`* steps`
+    with `full_length=True` — not a cheap one-step proxy, so building it
-    for flow) sequential forwards `sample.py` pays at inference, EVERY batch
+    costs the full `k_max` (`* steps` for flow) sequential forwards for every
+    row regardless of that row's own secondary count (`full_length=True`
+    disables `sample.py`'s inference-time row compaction — see that
+    function's docstring for why: the mixing below needs a real prediction
+    at every slot up to `k_max`, not just the valid ones). Paid EVERY batch
     this is called on (paid at train time too whenever teacher_forcing !=
     "always"). Fully detached: gradient only ever flows
     through the "real" target path each stage trainer already uses
@@ -306,7 +310,7 @@ def _assemble_stage2_ar_inputs_scheduled(
 
     was_training = model.training
     sec_cont_pred, sec_type_pred, _ = sample_secondaries_ar(
-        model, cond_cont, cond_cat, stage1_ctx, n_sec, steps=sample_steps
+        model, cond_cont, cond_cat, stage1_ctx, n_sec, steps=sample_steps, full_length=True
     )
     if was_training:
         model.train()

pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "giant"
-version = "0.3.17"
+version = "0.3.18"
 description = "Geant4 step-function surrogate via conditional flow matching"
 readme = "README.md"
 requires-python = ">=3.12"

tests/test_sample.py

@@ -341,6 +341,87 @@ def test_sample_secondaries_ar_none_n_sec_pred_without_stop_head_raises():
         sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
 
 
+# ── Row compaction: full_length=False (default, inference) must agree with
+# full_length=True (the pre-compaction behaviour, still exercised by
+# _assemble_stage2_ar_inputs_scheduled's training-time self-sample) ────────
+
+
+def _zero_randn(*size, **kwargs):
+    """Drop-in replacement for `torch.randn` that returns zeros of the same
+    shape — makes the ODE/WGAN noise deterministic so a compacted run and a
+    full_length run can be compared row-for-row regardless of how many
+    `torch.randn` calls each makes (compaction changes the batch size, and
+    therefore the RNG stream position, at every slot)."""
+    device = kwargs.get("device")
+    dtype = kwargs.get("dtype")
+    return torch.zeros(*size, device=device, dtype=dtype)
+
+
+@pytest.mark.parametrize("history", ["markov", "attention"])
+@pytest.mark.parametrize("generator", ["flow", "wgan"])
+def test_sample_secondaries_ar_compaction_matches_full_length_head_mode(generator, history, monkeypatch):
+    B, k_max, emb_dim = 4, 5, 6
+    decoder = _stage2_ar("physical", generator, emb_dim=emb_dim, k_max=k_max, history=history)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec_pred = torch.tensor([0, 1, 3, k_max])
+
+    monkeypatch.setattr(torch, "randn", _zero_randn)
+    sec_cont_c, sec_type_c, sec_valid_c = sample_secondaries_ar(
+        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2, full_length=False
+    )
+    sec_cont_f, sec_type_f, sec_valid_f = sample_secondaries_ar(
+        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2, full_length=True
+    )
+
+    assert torch.equal(sec_valid_c, sec_valid_f)
+    assert torch.equal(sec_valid_c, torch.arange(k_max).unsqueeze(0) < n_sec_pred.unsqueeze(1))
+    assert torch.allclose(sec_cont_c[sec_valid_c], sec_cont_f[sec_valid_f], atol=1e-4, rtol=1e-4)
+    assert torch.allclose(sec_type_c[sec_valid_c], sec_type_f[sec_valid_f], atol=1e-4, rtol=1e-4)
+
+
+@pytest.mark.parametrize("generator", ["flow", "wgan"])
+def test_sample_secondaries_ar_compaction_matches_full_length_stop_token(generator, monkeypatch):
+    """`n_sec_sampling="greedy"` keeps the stop decision itself deterministic
+    (no `torch.rand` draw), so only `torch.randn` needs zeroing."""
+    B, k_max = 6, 5
+    decoder = _stage2_ar_stop_token("physical", generator, n_sec_sampling="greedy", k_max=k_max)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+
+    monkeypatch.setattr(torch, "randn", _zero_randn)
+    sec_cont_c, sec_type_c, sec_valid_c = sample_secondaries_ar(
+        decoder, cond_cont, cond_cat, stage1_out, None, steps=2, full_length=False
+    )
+    sec_cont_f, sec_type_f, sec_valid_f = sample_secondaries_ar(
+        decoder, cond_cont, cond_cat, stage1_out, None, steps=2, full_length=True
+    )
+
+    assert torch.equal(sec_valid_c, sec_valid_f)
+    assert torch.allclose(sec_cont_c[sec_valid_c], sec_cont_f[sec_valid_f], atol=1e-4, rtol=1e-4)
+    assert torch.allclose(sec_type_c[sec_valid_c], sec_type_f[sec_valid_f], atol=1e-4, rtol=1e-4)
+
+
+def test_sample_secondaries_ar_full_length_ignores_n_sec_pred_zero_rows():
+    """A row with n_sec_pred == 0 would be dropped from the active set at
+    slot 0 under compaction (full_length=False) — full_length=True must
+    still run the model for it at every slot (only masked by sec_valid at
+    the end), matching _assemble_stage2_ar_inputs_scheduled's contract."""
+    B, k_max, emb_dim = 3, 4, 6
+    decoder = _stage2_ar("physical", "flow", emb_dim=emb_dim, k_max=k_max)
+    cond_cont, cond_cat = _cond(B)
+    stage1_out = torch.randn(B, X_DIM)
+    n_sec_pred = torch.tensor([0, 0, 0])
+    sec_cont, sec_type, sec_valid = sample_secondaries_ar(
+        decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2, full_length=True
+    )
+    assert not sec_valid.any()
+    # every slot still ran the model (not left at the zero-init default) —
+    # a real flow ODE output from randn-initialized noise is essentially
+    # never exactly zero.
+    assert not torch.allclose(sec_cont, torch.zeros_like(sec_cont))
+
+
 # ── resolve_n_sec: n_sec.mode = "head" sampling policy (gitea #86) ──────────
 
 

uv.lock

@@ -675,7 +675,7 @@ wheels = [
 
 [[package]]
 name = "giant"
-version = "0.3.17"
+version = "0.3.18"
 source = { editable = "." }
 dependencies = [
     { name = "numpy" },