chore: raise pyarrow ceiling to <26, bump to 25.0.1

25.0.0's changelog has no breaking changes for our usage (streaming parquet reads via polars/pyarrow) — the only Python-relevant deprecation is the `feather` module, which giant doesn't use. Tests, ruff, and ty all pass against 25.0.1. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TMdZFqXXig7i3XkirSUxef
chore: bump uv.lock and fix ruff 0.16 default-rule lint findings
2026-09-04 14:12:55 +02:00 · 2026-09-04 14:09:29 +02:00 · 2026-09-04 06:47:29 +00:00 · 2026-09-04 08:41:55 +02:00 · 2026-09-04 08:36:17 +02:00 · 2026-09-04 06:25:44 +00:00
71 changed files with 2355 additions and 1755 deletions
@@ -1,5 +1,5 @@
 [tool.bumpversion]
-current_version = "0.3.17"
+current_version = "0.3.20"
 parse = "(?P<major>\\d+)\\.(?P<minor>\\d+)\\.(?P<patch>\\d+)"
 serialize = ["{major}.{minor}.{patch}"]
 search = "{current_version}"
@@ -0,0 +1,20 @@
+name: Setup uv + sync deps
+description: >-
+  Install uv, point its cache at the runner-local mount, and sync the
+  project with the cpu + dev extras. Every CI job does this identically;
+  the caller must still mount /srv/act-runner-cache/uv:/uv-cache on its
+  own container (a composite action can't set that) and check out the
+  repo before this runs.
+
+runs:
+  using: composite
+  steps:
+    - uses: astral-sh/setup-uv@v5
+      with:
+        enable-cache: false
+    - shell: bash
+      run: |
+        echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
+        echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
+    - shell: bash
+      run: uv sync --extra cpu --extra dev
@@ -0,0 +1,63 @@
+#!/usr/bin/env bash
+# Build one complete release commit: version bump (or sync to a given
+# version, or no bump at all), changelog entry, and refreshed README
+# badges — all in a single commit, tagged at the end. Used by both the
+# merge-to-master release job and the hand-pushed-tag sync job, so "the
+# tag == a complete release commit" holds either way a release gets made.
+#
+# Usage: release-commit.sh [--no-bump|VERSION]
+#   (no argument)  bump the current patch version (normal merge path).
+#   --no-bump      don't touch the version — the branch already bumped it
+#                   (e.g. a manual minor/major bump); just build the
+#                   changelog/badge commit for whatever version is current.
+#   VERSION        sync the project to this exact version (tag-sync path).
+#
+# Preconditions: repo is checked out with full history (fetch-depth: 0),
+# `uv` is available and synced, and git user.name/user.email are configured.
+# Idempotent: if the resulting tag already exists, this is a no-op.
+set -euo pipefail
+
+VERSION_ARG="${1:-}"
+
+if [ "$VERSION_ARG" = "--no-bump" ]; then
+  echo "Version already bumped by this branch; using current version as-is"
+elif [ -n "$VERSION_ARG" ]; then
+  echo "Syncing project version to $VERSION_ARG"
+  uv version "$VERSION_ARG" --no-sync
+  uv lock
+else
+  CURRENT_VERSION=$(uv version --short)
+  echo "Bumping patch version from $CURRENT_VERSION"
+  uv run bump-my-version bump patch --current-version "$CURRENT_VERSION" --no-commit --no-tag
+  # bump-my-version's pre_commit_hooks (uv lock + git add uv.lock) only run
+  # on its own commit path, which we skipped with --no-commit — so do it here.
+  uv lock
+fi
+
+VERSION=$(uv version --short)
+TAG="v$VERSION"
+
+if git rev-parse "$TAG" >/dev/null 2>&1; then
+  echo "Tag $TAG already exists; nothing to do"
+  exit 0
+fi
+
+uv run git-cliff --tag "$TAG" --unreleased --prepend CHANGELOG.md
+
+TEST_COUNT=$(uv run pytest --collect-only -q 2>/dev/null | grep -oE '^[0-9]+ tests? collected' | grep -oE '^[0-9]+')
+sed -i -E "s|badge/version-[^-]+-informational|badge/version-${VERSION}-informational|" README.md
+sed -i -E "s|badge/tests-[0-9]+%20passing-brightgreen|badge/tests-${TEST_COUNT}%20passing-brightgreen|" README.md
+
+# .bumpversion.toml stores its own current_version, which bump-my-version
+# rewrites even with --no-commit — must be staged or the next run sees a
+# dirty tree and bump-my-version refuses (allow_dirty = false).
+git add pyproject.toml .bumpversion.toml uv.lock CHANGELOG.md README.md
+if git diff --cached --quiet; then
+  echo "Nothing changed; skipping release commit"
+  exit 0
+fi
+
+git commit -m "chore: release $TAG"
+git tag -a "$TAG" -m "$TAG"
+
+echo "Created release commit and tag $TAG"
@@ -12,6 +12,7 @@ env:
 jobs:
  ruff-check:
    name: Lint (ruff check)
+    if: ${{ !startsWith(github.ref, 'refs/tags/') }}
    runs-on: ubuntu-latest
    container:
      image: docker.gitea.com/runner-images:ubuntu-latest
@@ -19,17 +20,12 @@ jobs:
        - /srv/act-runner-cache/uv:/uv-cache
    steps:
      - uses: actions/checkout@v4
-      - 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
+      - uses: ./.gitea/actions/setup
      - run: uv run ruff check .

  ruff-format:
    name: Format (ruff format)
+    if: ${{ !startsWith(github.ref, 'refs/tags/') }}
    runs-on: ubuntu-latest
    container:
      image: docker.gitea.com/runner-images:ubuntu-latest
@@ -37,17 +33,12 @@ jobs:
        - /srv/act-runner-cache/uv:/uv-cache
    steps:
      - uses: actions/checkout@v4
-      - 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
+      - uses: ./.gitea/actions/setup
      - run: uv run ruff format --check .

  type-check:
    name: Type check (ty)
+    if: ${{ !startsWith(github.ref, 'refs/tags/') }}
    runs-on: ubuntu-latest
    container:
      image: docker.gitea.com/runner-images:ubuntu-latest
@@ -55,18 +46,13 @@ jobs:
        - /srv/act-runner-cache/uv:/uv-cache
    steps:
      - uses: actions/checkout@v4
-      - 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
+      - uses: ./.gitea/actions/setup
      - run: uv run ty check .

  test:
    name: Tests
-    needs: [ruff-check, type-check]
+    if: ${{ !startsWith(github.ref, 'refs/tags/') }}
+    needs: [ruff-check, ruff-format, type-check]
    runs-on: ubuntu-latest
    container:
      image: docker.gitea.com/runner-images:ubuntu-latest
@@ -74,21 +60,15 @@ jobs:
        - /srv/act-runner-cache/uv:/uv-cache
    steps:
      - uses: actions/checkout@v4
-      - 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
+      - uses: ./.gitea/actions/setup
      - run: uv run pytest --cov --cov-report=term-missing --cov-report=xml
-      - uses: actions/upload-artifact@v3
+      - uses: actions/upload-artifact@v4
        with:
          name: coverage-report
          path: coverage.xml

-  bump-version:
-    name: Bump version, tag, and update changelog on merge to master
+  release:
+    name: Release (bump, changelog, badges, tag) on merge to master
    needs: [ruff-check, ruff-format, type-check, test]
    if: github.ref == 'refs/heads/master' && github.event_name == 'push'
    runs-on: ubuntu-latest
@@ -98,7 +78,7 @@ jobs:
        - /srv/act-runner-cache/uv:/uv-cache
    steps:
      # CI_TOKEN needs write:repository scope (not just read) — this job
-      # pushes commits and tags to master, unlike ruff-check/ruff-format/
+      # pushes a commit and a tag to master, unlike ruff-check/ruff-format/
      # type-check/test above, which only need to check out the repo.
      - uses: actions/checkout@v4
        with:
@@ -114,22 +94,14 @@ jobs:
          else
            echo "is_merge=false" >> "$GITHUB_OUTPUT"
          fi
-      - uses: astral-sh/setup-uv@v5
-        if: steps.merge_check.outputs.is_merge == 'true'
-        with:
-          enable-cache: false
-      - run: |
-          echo "UV_CACHE_DIR=/uv-cache" >> "$GITHUB_ENV"
-          echo "UV_LINK_MODE=copy" >> "$GITHUB_ENV"
-        if: steps.merge_check.outputs.is_merge == 'true'
-      - run: uv sync --extra cpu --extra dev
+      - uses: ./.gitea/actions/setup
        if: steps.merge_check.outputs.is_merge == 'true'
      - name: Configure git identity
        if: steps.merge_check.outputs.is_merge == 'true'
        run: |
          git config user.name "gitea-actions"
          git config user.email "actions@git.larsbogner.de"
-      - name: Bump patch version if this merge didn't already bump it
+      - name: Build the release commit
        if: steps.merge_check.outputs.is_merge == 'true'
        run: |
          OLD_VERSION=$(git show "${{ github.event.before }}:pyproject.toml" 2>/dev/null | grep -m1 '^version = ' | sed -E 's/version = "(.*)"/\1/')
@@ -140,70 +112,64 @@ jobs:
          fi
          if [ "$OLD_VERSION" = "$CURRENT_VERSION" ]; then
            echo "Version unchanged by this merge ($CURRENT_VERSION); bumping patch"
-            uv run bump-my-version bump patch --current-version "$CURRENT_VERSION"
+            .gitea/scripts/release-commit.sh
          else
-            echo "Branch already bumped the version ($OLD_VERSION -> $CURRENT_VERSION); skipping auto-bump"
+            echo "Branch already bumped the version ($OLD_VERSION -> $CURRENT_VERSION); building release commit without bumping"
+            .gitea/scripts/release-commit.sh --no-bump
          fi
-      - name: Update changelog for the current version if not already tagged
+      - name: Push the release commit and its tag together
        if: steps.merge_check.outputs.is_merge == 'true'
        run: |
          VERSION=$(uv version --short)
          TAG="v$VERSION"
-          if git rev-parse "$TAG" >/dev/null 2>&1; then
-            echo "Tag $TAG already exists; skipping changelog update"
+          if git rev-parse "$TAG" >/dev/null 2>&1 && [ "$(git rev-parse "$TAG^{commit}")" = "$(git rev-parse HEAD)" ]; then
+            git push --atomic origin HEAD:master "refs/tags/$TAG"
          else
-            uv run git-cliff --tag "$TAG" --unreleased --prepend CHANGELOG.md
-            git add CHANGELOG.md
-            if ! git diff --cached --quiet -- CHANGELOG.md; then
-              git commit -m "chore: update changelog for $TAG"
-            else
-              git restore --staged CHANGELOG.md
-            fi
-          fi
-      - name: Push commits and tag the current version
-        if: steps.merge_check.outputs.is_merge == 'true'
-        run: |
-          git push origin HEAD:master
-          VERSION=$(uv version --short)
-          TAG="v$VERSION"
-          if git rev-parse "$TAG" >/dev/null 2>&1; then
-            echo "Tag $TAG already exists"
-          else
-            git tag -a "$TAG" -m "$TAG"
-            git push origin "refs/tags/$TAG"
+            echo "No new release commit/tag to push (already released, or nothing changed)"
+            git push origin HEAD:master
          fi

  sync-version-on-tag:
-    name: Sync project version with tag
+    name: Sync project version with tag (hand-pushed tags only)
    if: startsWith(github.ref, 'refs/tags/')
    runs-on: ubuntu-latest
+    container:
+      image: docker.gitea.com/runner-images:ubuntu-latest
+      volumes:
+        - /srv/act-runner-cache/uv:/uv-cache
    steps:
      - uses: actions/checkout@v4
        with:
          token: ${{ secrets.CI_TOKEN }}
-      - uses: astral-sh/setup-uv@v5
-      - name: Check tag against project version, update if they differ
+          fetch-depth: 0
+      - name: Require the tagged commit to already be on master
+        run: |
+          git fetch origin master
+          if ! git merge-base --is-ancestor "${{ github.sha }}" origin/master; then
+            echo "::error::Tag ${GITHUB_REF_NAME} points at a commit not on master; refusing to publish an unreviewed tree. Push the commit to master first, or delete and re-push the tag once it is."
+            exit 1
+          fi
+      - uses: ./.gitea/actions/setup
+      - name: Check tag against project version, build a release commit if they differ
        run: |
          TAG_VERSION="${GITHUB_REF_NAME#v}"
          CURRENT_VERSION=$(uv version --short)
-          if [ "$TAG_VERSION" != "$CURRENT_VERSION" ]; then
-            echo "Tag version ($TAG_VERSION) != project version ($CURRENT_VERSION); updating pyproject.toml"
-            uv version "$TAG_VERSION" --no-sync
+          if [ "$TAG_VERSION" = "$CURRENT_VERSION" ]; then
+            echo "Tag version matches project version ($CURRENT_VERSION)"
+          else
+            echo "Tag version ($TAG_VERSION) != project version ($CURRENT_VERSION); building a release commit"
            git config user.name "gitea-actions"
            git config user.email "actions@git.larsbogner.de"
-            git add pyproject.toml uv.lock
-            git commit -m "chore: sync project version to tag ${GITHUB_REF_NAME}"
+            .gitea/scripts/release-commit.sh "$TAG_VERSION"
            git push origin HEAD:master
            git push origin ":refs/tags/${GITHUB_REF_NAME}"
            git tag -f "${GITHUB_REF_NAME}" HEAD
            git push origin "refs/tags/${GITHUB_REF_NAME}"
-          else
-            echo "Tag version matches project version ($CURRENT_VERSION)"
          fi

  publish-package:
    name: Publish package to Gitea package registry
-    needs: [ruff-check, ruff-format, type-check, test, sync-version-on-tag]
+    needs: [sync-version-on-tag]
    if: startsWith(github.ref, 'refs/tags/')
    runs-on: ubuntu-latest
    container:
@@ -216,12 +182,7 @@ jobs:
      - uses: actions/checkout@v4
        with:
          ref: ${{ github.ref_name }}
-      - 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"
+      - uses: ./.gitea/actions/setup
      - run: uv build
      # CI_TOKEN needs write:package scope (in addition to write:repository,
      # used elsewhere) for this upload to authenticate.
@@ -1,5 +1,27 @@
 # Changelog

+## [0.3.20] - 2026-09-04
+
+### Changed
+
+- Fix: dereference annotated tag to its commit before the push comparison
+
+## [0.3.19] - 2026-09-04
+
+### Changed
+
+- Chore: update README badges (version 0.3.18, 1138 tests)
+
+- Ci: restructure release pipeline into a single atomic release commit
+
+## [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
@@ -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.20-informational)](CHANGELOG.md)
+[![tests](https://img.shields.io/badge/tests-1138%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)
-
-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
+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 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 |
-| 3–5 | `post_dir` in local frame | unit vector |
-| 6–8 | `travel_dir` (`post_pos − pre_pos`) in local frame | unit vector |
+| 1–2 | `edep_logit`, `sec_logit` | ALR coordinates of the deposit / secondary / post-energy simplex |
+| 3–5 | `post_dir` | unit vector, local frame (`pre_dir = ẑ`) |
+| 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.
- `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)`.
+Energy logits decode via `softmax([edep_logit, sec_logit, 0]) × pre_E`, so
+`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)
- `one_shot` — all `K_MAX` slots generated in a single forward pass, masked past the predicted `n_sec`
+**Conditioning (15D).** 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 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.
- **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`.
- **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`).
- **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.
- **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.
- Train/val split is by `event_id` (`--seed`-controlled), not row shuffle, so correlated steps from the same shower never leak across the split.
- 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.
+Every default lives in one place: frozen dataclasses in `giant/config.py`,
+composed into `GiantConfig` (`conditioning` / `stage1_model` / `stage2_model` /
+`train`). `DEFAULT_CONFIG` is *generated* from `GiantConfig().to_dict()` rather
+than hand-maintained, so the dataclasses can't drift from what actually gets
+merged. TOML config keys are validated against that shape — an unknown key is
+rejected with a did-you-mean suggestion. Precedence: CLI flag > `--config` file
+> default.

-## Project structure
+```toml
+# config.toml — resolved shape of the four blocks
+[conditioning]
+particle.type = "physical"
+material.type = "physical"

-```
-giant/
-├── giant/
-│   ├── data/
-│   │   ├── loader.py       # parquet → numpy arrays (incl. streaming/chunked reads)
-│   │   ├── transforms.py   # log transforms, local-frame rotation, energy simplex, secondary encode/decode
-│   │   ├── dataset.py      # StepsDataset / StreamingStepsDataset (PyTorch)
-│   │   └── setup_cache.py  # sidecar cache for the pre-epoch setup scan (vocab/split/normalizers)
-│   ├── model/
-│   │   ├── models.py       # Stage1Model, Stage2OneShot, Stage2Autoregressive, CriticModel
-│   │   ├── builders.py     # build_models / build_critics — config dict → assembled stage models
-│   │   ├── 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/
+[stage1_model]
+generator = "flow"
+
+[stage2_model]
+generator = "wgan"
+decoder   = "autoregressive"
+
+[train]
+epochs     = 100
+batch_size = 4096
+lr         = 3e-4
 ```

-## 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       # 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`)
+uv sync --extra cpu --extra dev     # everything needed to develop
 ```

-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
-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 pytest            # 964 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.
@@ -37,7 +37,7 @@ commit_preprocessors = [
 protect_breaking_commits = false
 commit_parsers = [
    { message = "^Merge ", skip = true },
-    { message = "^chore: (bump version|update changelog|sync project version)", skip = true },
+    { message = "^chore: (release|bump version|update changelog|sync project version)", skip = true },
    { message = "^Add", group = "<!-- 0 -->Added" },
    { message = "^(Fix|Clamp|Clip)", group = "<!-- 1 -->Fixed" },
    { message = "^(Remove|Drop|Deprecate)", group = "<!-- 2 -->Removed" },
@@ -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:
-# sample.sample_secondaries_ar loops `for k in range(k_max)` unconditionally
-# — 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:
+# COST, measured (RTX 4070, bs 4096, 10 ODE steps), not estimated — but see
+# the row-compaction note below, which changes the INFERENCE side of this:
 #   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
@@ -32,27 +32,27 @@ from giant.analysis.runtime_estimate import RUNTIME_SAFETY_MARGIN, estimate_runt
 from giant.analysis.sources import RolloutSpec, Side

 __all__ = [
-    "build_catalog",
-    "catalog_ids",
-    "get_spec",
+    "RUNTIME_SAFETY_MARGIN",
+    "Context",
    "LoadedRollout",
+    "Partial",
+    "Reduced",
+    "RolloutSpec",
    "RunMeta",
+    "Side",
    "SubmitConfig",
+    "build_catalog",
+    "build_context",
+    "catalog_ids",
    "compute_one",
    "compute_reduced",
    "derive_run_dir",
+    "estimate_runtime_s",
+    "get_spec",
    "load_rollout_yaml",
    "load_rollout_yamls",
    "merge_all",
    "merge_one",
    "prep",
    "write_submit",
-    "Context",
-    "build_context",
-    "Partial",
-    "Reduced",
-    "RolloutSpec",
-    "Side",
-    "RUNTIME_SAFETY_MARGIN",
-    "estimate_runtime_s",
 ]
@@ -38,8 +38,8 @@ variable x grouping, secondaries, ...) into concrete specs.

 from __future__ import annotations

+from collections.abc import Callable
 from dataclasses import asdict, dataclass
-from typing import Callable

 import numpy as np
 import polars as pl
@@ -101,7 +101,7 @@ class Bundle:
        reference,
        ctx: Context,
        chunk: tuple[int, int] | None = None,
-    ) -> "Bundle":
+    ) -> Bundle:
        """Open the reference + every rollout, optionally restricted to one event-disjoint chunk.

        ``chunk = (chunk_index, n_chunks)`` filters every side to
@@ -228,7 +228,7 @@ class RunMeta:
        Path(path).write_text(json.dumps(self.__dict__, indent=2))

    @classmethod
-    def load(cls, path: str | Path) -> "RunMeta":
+    def load(cls, path: str | Path) -> RunMeta:
        return cls(**json.loads(Path(path).read_text()))


@@ -50,7 +50,7 @@ class Context:
        Path(path).write_text(json.dumps(asdict(self), indent=2))

    @classmethod
-    def load(cls, path: str | Path) -> "Context":
+    def load(cls, path: str | Path) -> Context:
        d = json.loads(Path(path).read_text())
        d["var_ranges"] = {k: tuple(v) for k, v in d["var_ranges"].items()}
        d["sec_energy_range"] = tuple(d["sec_energy_range"])
@@ -46,7 +46,7 @@ def pdg_label(code: int) -> str:

 def material_label(name: str) -> str:
    """Display label for a Geant4 material, dropping the ``G4_`` prefix."""
-    return name[3:] if name.startswith("G4_") else name
+    return name.removeprefix("G4_")


 def energy_bin_edges(incident_E: np.ndarray, n_bins: int = 4) -> np.ndarray:
@@ -46,7 +46,7 @@ class Reduced:
        Path(path).write_text(json.dumps(asdict(self)))

    @classmethod
-    def load(cls, path: str | Path) -> "Reduced":
+    def load(cls, path: str | Path) -> Reduced:
        return cls(**json.loads(Path(path).read_text()))


@@ -70,5 +70,5 @@ class Partial:
        Path(path).write_text(json.dumps(asdict(self)))

    @classmethod
-    def load(cls, path: str | Path) -> "Partial":
+    def load(cls, path: str | Path) -> Partial:
        return cls(**json.loads(Path(path).read_text()))
@@ -27,8 +27,8 @@ from pathlib import Path

 import numpy as np
 import plotstyle as ps
-from matplotlib.colors import LogNorm
 import yaml
+from matplotlib.colors import LogNorm

 from giant.analysis.reduced import Reduced

@@ -58,10 +58,10 @@ _COLS = (

@dataclass
 class _RouterHandle:
-    router: "torch.nn.Module"
+    router: torch.nn.Module
    pdg_map: dict[int, int]
    mat_map: dict[str, int]
-    cond_normalizer: "Normalizer"
+    cond_normalizer: Normalizer
    particle_conditioning: str
    material_conditioning: str
    router_type: str
@@ -233,7 +233,7 @@ def _gating_entry(checkpoint: str | Path | None, r_phys: pl.LazyFrame, t_phys: p
    return {"router_type": handle.router_type, "n_experts": handle.router.n_experts, **sides}


-def compute_router_gating(rollouts: dict[str, "RolloutSide"], t_phys: pl.LazyFrame, seed: int = 0) -> Reduced:
+def compute_router_gating(rollouts: dict[str, RolloutSide], t_phys: pl.LazyFrame, seed: int = 0) -> Reduced:
    """`Reduced` for the router-gating figure: one panel-pair per rollout with
    an enabled MoE router, or an explanatory note if none of them have one."""
    series = {}
@@ -289,7 +289,7 @@ def _specialization_entry(
    }


-def compute_router_specialization(rollouts: dict[str, "RolloutSide"], t_phys: pl.LazyFrame, seed: int = 0) -> Reduced:
+def compute_router_specialization(rollouts: dict[str, RolloutSide], t_phys: pl.LazyFrame, seed: int = 0) -> Reduced:
    """`Reduced` for the router-specialization figure, one curve per rollout with
    an enabled MoE router (see `_specialization_entry`)."""
    series = {}
@@ -331,7 +331,7 @@ def _share_by_pdg_entry(


 def compute_router_share_by_pdg(
-    rollouts: dict[str, "RolloutSide"], t_phys: pl.LazyFrame, top_pdgs: list[int], seed: int = 0
+    rollouts: dict[str, RolloutSide], t_phys: pl.LazyFrame, top_pdgs: list[int], seed: int = 0
 ) -> Reduced:
    """`Reduced` for the router expert-share-by-species figure, one panel-pair
    per rollout with an enabled MoE router."""
@@ -375,7 +375,7 @@ def _share_by_process_entry(checkpoint: str | Path | None, t_phys: pl.LazyFrame,


 def compute_router_share_by_process(
-    rollouts: dict[str, "RolloutSide"], t_phys: pl.LazyFrame, seed: int = 0, top_k: int = _TOP_K_PROCESS
+    rollouts: dict[str, RolloutSide], t_phys: pl.LazyFrame, seed: int = 0, top_k: int = _TOP_K_PROCESS
 ) -> Reduced:
    """Stacked-bar share of each physics process dispatched to each expert, one
    panel per rollout checkpoint with an enabled MoE router.
@@ -35,7 +35,7 @@ _NOTE_NOT_APPLICABLE = (
 )


-def compute_type_embedding_l1_distance(rollouts: dict[str, "RolloutSide"]) -> Reduced:
+def compute_type_embedding_l1_distance(rollouts: dict[str, RolloutSide]) -> Reduced:
    """`Reduced` for the type-embedding-distance figure: one series per rollout
    whose checkpoint populated the diagnostic, or an explanatory note if none did.

@@ -1,16 +1,15 @@
 from __future__ import annotations

-from collections import Counter
-from datetime import datetime, timezone
-from enum import Enum
 import math
-from pathlib import Path
 import re
-from typing import TYPE_CHECKING, Optional, cast
 import uuid as uuid_mod
+from collections import Counter
+from datetime import UTC, datetime
+from enum import Enum
+from pathlib import Path
+from typing import TYPE_CHECKING, Annotated, cast

 import typer
-from typing_extensions import Annotated

 if TYPE_CHECKING:
    import numpy as np
@@ -120,7 +119,7 @@ def _parse_router_axis_flags(specs: list[str]) -> dict[str, object]:
    return out


-def _parse_set_flags(specs: Optional[list[str]]) -> dict[str, object]:
+def _parse_set_flags(specs: list[str] | None) -> dict[str, object]:
    """Parse repeated `--set dotted.path=value` flags into a dict, typing
    each value with `_coerce_scalar` the same way a TOML file's native types
    would arrive. Validation against the inference-safe allowlist happens
@@ -196,7 +195,7 @@ def _write_prediction_ref(
        "output": str(out),
        "dataset": str(dataset_path),
        "checkpoint": str(checkpoint.resolve()),
-        "timestamp": datetime.now(timezone.utc).isoformat(),
+        "timestamp": datetime.now(UTC).isoformat(),
    }
    if comment is not None:
        ref["comment"] = comment
@@ -285,52 +284,52 @@ class Weights(str, Enum):
 def train(
    data: Annotated[Path, typer.Argument(help="Parquet file or directory of parquet files")],
    config: Annotated[
-        Optional[Path],
+        Path | None,
        typer.Option("--config", "-c", help="TOML config file (overridden by explicit flags)"),
    ] = None,
    mode: Annotated[
-        Optional[Mode],
+        Mode | None,
        typer.Option("--mode", "-m", help="Generative model: flow matching or DDPM"),
    ] = None,
-    epochs: Annotated[Optional[int], typer.Option("--epochs", "-e")] = None,
+    epochs: Annotated[int | None, typer.Option("--epochs", "-e")] = None,
    batch_size: Annotated[
-        Optional[str],
+        str | None,
        typer.Option(
            "--batch-size",
            "-b",
            help="Integer, or 'auto' to estimate from free GPU memory (cuda devices only)",
        ),
    ] = None,
-    lr: Annotated[Optional[float], typer.Option("--lr", "-l")] = None,
+    lr: Annotated[float | None, typer.Option("--lr", "-l")] = None,
    weight_decay: Annotated[
-        Optional[float],
+        float | None,
        typer.Option("--weight-decay", "-W", help="AdamW weight decay (default: 0.01)"),
    ] = None,
    ema_decay: Annotated[
-        Optional[float],
+        float | None,
        typer.Option(
            "--ema-decay",
            help="EMA decay for a shadow copy of the model weights, saved "
            "alongside the raw weights in checkpoints (0 disables; default: 0.9999)",
        ),
    ] = None,
-    warmup_epochs: Annotated[Optional[int], typer.Option("--warmup-epochs", "-w")] = None,
-    hidden_dim: Annotated[Optional[int], typer.Option("--hidden-dim", "-H")] = None,
-    n_blocks: Annotated[Optional[int], typer.Option("--n-blocks", "-n")] = None,
-    emb_dim: Annotated[Optional[int], typer.Option("--emb-dim", "-E")] = None,
+    warmup_epochs: Annotated[int | None, typer.Option("--warmup-epochs", "-w")] = None,
+    hidden_dim: Annotated[int | None, typer.Option("--hidden-dim", "-H")] = None,
+    n_blocks: Annotated[int | None, typer.Option("--n-blocks", "-n")] = None,
+    emb_dim: Annotated[int | None, typer.Option("--emb-dim", "-E")] = None,
    dropout: Annotated[
-        Optional[float],
+        float | None,
        typer.Option("--dropout", "-d", help="Dropout probability in ResBlocks (default: 0.1)"),
    ] = None,
    stage1_generator: Annotated[
-        Optional[Mode],
+        Mode | None,
        typer.Option(
            "--stage1-generator",
            help="Stage 1's generative objective — overrides --mode for stage 1 only",
        ),
    ] = None,
    stage1_hidden_dim: Annotated[
-        Optional[int],
+        int | None,
        typer.Option(
            "--stage1-hidden-dim",
            help="Overrides --hidden-dim for stage 1 only (same effect today; "
@@ -339,15 +338,15 @@ def train(
        ),
    ] = None,
    stage1_n_res_blocks: Annotated[
-        Optional[int],
+        int | None,
        typer.Option("--stage1-n-res-blocks", help="Overrides --n-blocks for stage 1 only"),
    ] = None,
    stage1_dropout: Annotated[
-        Optional[float],
+        float | None,
        typer.Option("--stage1-dropout", help="Overrides --dropout for stage 1 only"),
    ] = None,
    stage2_generator: Annotated[
-        Optional[Mode],
+        Mode | None,
        typer.Option(
            "--stage2-generator",
            help="Stage 2's generative objective — overrides --mode for stage 2 "
@@ -356,19 +355,19 @@ def train(
        ),
    ] = None,
    stage2_hidden_dim: Annotated[
-        Optional[int],
+        int | None,
        typer.Option("--stage2-hidden-dim", help="Stage 2 trunk width"),
    ] = None,
    stage2_n_res_blocks: Annotated[
-        Optional[int],
+        int | None,
        typer.Option("--stage2-n-res-blocks", help="Stage 2 trunk depth"),
    ] = None,
    stage2_dropout: Annotated[
-        Optional[float],
+        float | None,
        typer.Option("--stage2-dropout", help="Dropout inside stage 2's ResBlocks"),
    ] = None,
    stage2_decoder: Annotated[
-        Optional[Decoder],
+        Decoder | None,
        typer.Option(
            "--stage2-decoder",
            help="one_shot: predict all k_max secondary slots at once (v0.2 "
@@ -377,7 +376,7 @@ def train(
        ),
    ] = None,
    stage2_k_max: Annotated[
-        Optional[int],
+        int | None,
        typer.Option(
            "--stage2-k-max",
            help="Maximum secondary slots (fixed width under one_shot, a "
@@ -385,14 +384,14 @@ def train(
        ),
    ] = None,
    stage2_context_dim: Annotated[
-        Optional[int],
+        int | None,
        typer.Option(
            "--stage2-context-dim",
            help="Width of the projected stage-1 outcome fed into stage 2's conditioning (default: 64)",
        ),
    ] = None,
    stage2_stage1_context: Annotated[
-        Optional[Stage1Context],
+        Stage1Context | None,
        typer.Option(
            "--stage2-stage1-context",
            help="What stage 2 conditions on during training: 'truth' (the "
@@ -402,7 +401,7 @@ def train(
        ),
    ] = None,
    conditioning: Annotated[
-        Optional[Conditioning],
+        Conditioning | None,
        typer.Option(
            "--conditioning",
            help="Input conditioning: continuous physical properties "
@@ -411,7 +410,7 @@ def train(
        ),
    ] = None,
    router: Annotated[
-        Optional[bool],
+        bool | None,
        typer.Option(
            "--router/--no-router",
            help="Route both stages through a mixture of small experts "
@@ -419,12 +418,12 @@ def train(
        ),
    ] = None,
    router_type: Annotated[
-        Optional[str],
+        str | None,
        typer.Option("--router-type", help="Router implementation name (see ROUTER_REGISTRY)"),
    ] = None,
-    n_experts: Annotated[Optional[int], typer.Option("--n-experts", help="Number of routed experts")] = None,
+    n_experts: Annotated[int | None, typer.Option("--n-experts", help="Number of routed experts")] = None,
    router_axis: Annotated[
-        Optional[list[str]],
+        list[str] | None,
        typer.Option(
            "--router-axis",
            help="Composed-router axis spec 'type:key=val,key=val' (repeatable; "
@@ -434,59 +433,59 @@ def train(
        ),
    ] = None,
    n_critic: Annotated[
-        Optional[int],
+        int | None,
        typer.Option(
            "--n-critic",
            help="WGAN-GP (--mode wgan only): critic updates per generator update (default: 5)",
        ),
    ] = None,
    gp_weight: Annotated[
-        Optional[float],
+        float | None,
        typer.Option(
            "--gp-weight",
            help="WGAN-GP (--mode wgan only): gradient-penalty coefficient (default: 10.0)",
        ),
    ] = None,
    noise_dim: Annotated[
-        Optional[int],
+        int | None,
        typer.Option(
            "--noise-dim",
            help="WGAN (--mode wgan only): generator input noise-vector width (default: 64)",
        ),
    ] = None,
    critic_lr: Annotated[
-        Optional[float],
+        float | None,
        typer.Option(
            "--critic-lr",
            help="WGAN-GP (--mode wgan only): critic learning rate (default: same as --lr)",
        ),
    ] = None,
    stage1_n_critic: Annotated[
-        Optional[int],
+        int | None,
        typer.Option("--stage1-n-critic", help="Overrides --n-critic for stage 1 only"),
    ] = None,
    stage1_gp_weight: Annotated[
-        Optional[float],
+        float | None,
        typer.Option("--stage1-gp-weight", help="Overrides --gp-weight for stage 1 only"),
    ] = None,
    stage1_noise_dim: Annotated[
-        Optional[int],
+        int | None,
        typer.Option("--stage1-noise-dim", help="Overrides --noise-dim for stage 1 only"),
    ] = None,
    stage1_critic_lr: Annotated[
-        Optional[float],
+        float | None,
        typer.Option("--stage1-critic-lr", help="Overrides --critic-lr for stage 1 only"),
    ] = None,
    stage2_n_critic: Annotated[
-        Optional[int],
+        int | None,
        typer.Option("--stage2-n-critic", help="Overrides --n-critic for stage 2 only"),
    ] = None,
    stage2_gp_weight: Annotated[
-        Optional[float],
+        float | None,
        typer.Option("--stage2-gp-weight", help="Overrides --gp-weight for stage 2 only"),
    ] = None,
    stage2_noise_dim: Annotated[
-        Optional[int],
+        int | None,
        typer.Option(
            "--stage2-noise-dim",
            help="Overrides --noise-dim for stage 2 only; under "
@@ -494,39 +493,39 @@ def train(
        ),
    ] = None,
    stage2_critic_lr: Annotated[
-        Optional[float],
+        float | None,
        typer.Option("--stage2-critic-lr", help="Overrides --critic-lr for stage 2 only"),
    ] = None,
    stage1_critic_hidden_dim: Annotated[
-        Optional[int],
+        int | None,
        typer.Option(
            "--stage1-critic-hidden-dim",
            help="WGAN-GP (--mode wgan only): critic width for stage 1 (default: same as generator's hidden_dim)",
        ),
    ] = None,
    stage1_critic_n_res_blocks: Annotated[
-        Optional[int],
+        int | None,
        typer.Option(
            "--stage1-critic-n-res-blocks",
            help="WGAN-GP (--mode wgan only): critic depth for stage 1 (default: same as generator's n_res_blocks)",
        ),
    ] = None,
    stage2_critic_hidden_dim: Annotated[
-        Optional[int],
+        int | None,
        typer.Option(
            "--stage2-critic-hidden-dim",
            help="WGAN-GP (--mode wgan only): critic width for stage 2 (default: same as generator's hidden_dim)",
        ),
    ] = None,
    stage2_critic_n_res_blocks: Annotated[
-        Optional[int],
+        int | None,
        typer.Option(
            "--stage2-critic-n-res-blocks",
            help="WGAN-GP (--mode wgan only): critic depth for stage 2 (default: same as generator's n_res_blocks)",
        ),
    ] = None,
    stage1_init_from: Annotated[
-        Optional[Path],
+        Path | None,
        typer.Option(
            "--stage1-init-from",
            help="Checkpoint .pt to load stage 1's weights from before training starts "
@@ -535,33 +534,33 @@ def train(
        ),
    ] = None,
    stage1_freeze: Annotated[
-        Optional[bool],
+        bool | None,
        typer.Option(
            "--stage1-freeze/--no-stage1-freeze",
            help="Never update stage 1's weights (requires --stage1-init-from, or --resume)",
        ),
    ] = None,
    stage2_init_from: Annotated[
-        Optional[Path],
+        Path | None,
        typer.Option(
            "--stage2-init-from",
            help="Checkpoint .pt to load stage 2's weights from before training starts (gitea #42)",
        ),
    ] = None,
    stage2_freeze: Annotated[
-        Optional[bool],
+        bool | None,
        typer.Option(
            "--stage2-freeze/--no-stage2-freeze",
            help="Never update stage 2's weights (requires --stage2-init-from, or --resume)",
        ),
    ] = None,
-    val_fraction: Annotated[Optional[float], typer.Option("--val-fraction", "-f")] = None,
+    val_fraction: Annotated[float | None, typer.Option("--val-fraction", "-f")] = None,
    seed: Annotated[
-        Optional[int],
+        int | None,
        typer.Option("--seed", "-s", help="Random seed for reproducibility"),
    ] = None,
    validate_every: Annotated[
-        Optional[int],
+        int | None,
        typer.Option(
            "--validate-every",
            "-v",
@@ -569,7 +568,7 @@ def train(
        ),
    ] = None,
    validate_steps: Annotated[
-        Optional[int],
+        int | None,
        typer.Option(
            "--validate-steps",
            "-t",
@@ -578,7 +577,7 @@ def train(
        ),
    ] = None,
    max_val_batches: Annotated[
-        Optional[int],
+        int | None,
        typer.Option(
            "--max-val-batches",
            help="Cap the per-epoch val-loss pass to N batches (0 = full val set every epoch; default: 200)",
@@ -609,7 +608,7 @@ def train(
        ),
    ] = False,
    out: Annotated[
-        Optional[Path],
+        Path | None,
        typer.Option(
            "--out",
            "-o",
@@ -618,31 +617,31 @@ def train(
        ),
    ] = None,
    device: Annotated[
-        Optional[str],
+        str | None,
        typer.Option("--device", "-D", help="cpu | cuda | mps (default: auto)"),
    ] = None,
-    num_workers: Annotated[Optional[int], typer.Option("--num-workers", "-j")] = None,
+    num_workers: Annotated[int | None, typer.Option("--num-workers", "-j")] = None,
    resume: Annotated[
-        Optional[Path],
+        Path | None,
        typer.Option("--resume", "-r", help="Checkpoint .pt to resume training from"),
    ] = None,
    wandb: Annotated[
-        Optional[bool],
+        bool | None,
        typer.Option(
            "--wandb/--no-wandb",
            help="Log per-epoch training metrics to Weights & Biases (requires `uv sync --extra wandb`)",
        ),
    ] = None,
    wandb_project: Annotated[
-        Optional[str],
+        str | None,
        typer.Option("--wandb-project", help="W&B project name (default: giant)"),
    ] = None,
    wandb_run_name: Annotated[
-        Optional[str],
+        str | None,
        typer.Option("--wandb-run-name", help="W&B run name (default: out_dir name)"),
    ] = None,
    wandb_log_every: Annotated[
-        Optional[int],
+        int | None,
        typer.Option(
            "--wandb-log-every",
            help="Log batch-level loss/grad_norm/lr to W&B every N optimizer "
@@ -650,7 +649,7 @@ def train(
        ),
    ] = None,
    precision: Annotated[
-        Optional[str],
+        str | None,
        typer.Option(
            "--precision",
            help="Training-step autocast precision: 'fp32' (default) or "
@@ -664,7 +663,7 @@ def train(
    from giant.pipeline import run_train_job

    batch_size_auto = False
-    batch_size_value: Optional[int] = None
+    batch_size_value: int | None = None
    if batch_size is not None:
        if batch_size.strip().lower() == "auto":
            batch_size_auto = True
@@ -794,52 +793,52 @@ def train(
@app.command("new-run")
 def new_run(
    config: Annotated[
-        Optional[Path],
+        Path | None,
        typer.Option(
            "--config",
            "-c",
            help="Base TOML to start from (default: built-in defaults)",
        ),
    ] = None,
-    mode: Annotated[Optional[Mode], typer.Option("--mode", "-m")] = None,
-    epochs: Annotated[Optional[int], typer.Option("--epochs", "-e")] = None,
-    batch_size: Annotated[Optional[int], typer.Option("--batch-size", "-b")] = None,
-    lr: Annotated[Optional[float], typer.Option("--lr", "-l")] = None,
-    hidden_dim: Annotated[Optional[int], typer.Option("--hidden-dim", "-H")] = None,
-    n_blocks: Annotated[Optional[int], typer.Option("--n-blocks", "-n")] = None,
-    emb_dim: Annotated[Optional[int], typer.Option("--emb-dim", "-E")] = None,
-    dropout: Annotated[Optional[float], typer.Option("--dropout", "-d")] = None,
-    stage1_generator: Annotated[Optional[Mode], typer.Option("--stage1-generator")] = None,
-    stage1_hidden_dim: Annotated[Optional[int], typer.Option("--stage1-hidden-dim")] = None,
-    stage1_n_res_blocks: Annotated[Optional[int], typer.Option("--stage1-n-res-blocks")] = None,
-    stage1_dropout: Annotated[Optional[float], typer.Option("--stage1-dropout")] = None,
-    stage2_generator: Annotated[Optional[Mode], typer.Option("--stage2-generator")] = None,
-    stage2_hidden_dim: Annotated[Optional[int], typer.Option("--stage2-hidden-dim")] = None,
-    stage2_n_res_blocks: Annotated[Optional[int], typer.Option("--stage2-n-res-blocks")] = None,
-    stage2_dropout: Annotated[Optional[float], typer.Option("--stage2-dropout")] = None,
-    stage2_decoder: Annotated[Optional[Decoder], typer.Option("--stage2-decoder")] = None,
-    stage2_k_max: Annotated[Optional[int], typer.Option("--stage2-k-max")] = None,
-    stage2_context_dim: Annotated[Optional[int], typer.Option("--stage2-context-dim")] = None,
-    stage2_stage1_context: Annotated[Optional[Stage1Context], typer.Option("--stage2-stage1-context")] = None,
-    stage1_init_from: Annotated[Optional[Path], typer.Option("--stage1-init-from")] = None,
-    stage1_freeze: Annotated[Optional[bool], typer.Option("--stage1-freeze/--no-stage1-freeze")] = None,
-    stage2_init_from: Annotated[Optional[Path], typer.Option("--stage2-init-from")] = None,
-    stage2_freeze: Annotated[Optional[bool], typer.Option("--stage2-freeze/--no-stage2-freeze")] = None,
-    conditioning: Annotated[Optional[Conditioning], typer.Option("--conditioning")] = None,
-    router: Annotated[Optional[bool], typer.Option("--router/--no-router")] = None,
-    router_type: Annotated[Optional[str], typer.Option("--router-type")] = None,
-    n_experts: Annotated[Optional[int], typer.Option("--n-experts")] = None,
-    router_axis: Annotated[Optional[list[str]], typer.Option("--router-axis")] = None,
+    mode: Annotated[Mode | None, typer.Option("--mode", "-m")] = None,
+    epochs: Annotated[int | None, typer.Option("--epochs", "-e")] = None,
+    batch_size: Annotated[int | None, typer.Option("--batch-size", "-b")] = None,
+    lr: Annotated[float | None, typer.Option("--lr", "-l")] = None,
+    hidden_dim: Annotated[int | None, typer.Option("--hidden-dim", "-H")] = None,
+    n_blocks: Annotated[int | None, typer.Option("--n-blocks", "-n")] = None,
+    emb_dim: Annotated[int | None, typer.Option("--emb-dim", "-E")] = None,
+    dropout: Annotated[float | None, typer.Option("--dropout", "-d")] = None,
+    stage1_generator: Annotated[Mode | None, typer.Option("--stage1-generator")] = None,
+    stage1_hidden_dim: Annotated[int | None, typer.Option("--stage1-hidden-dim")] = None,
+    stage1_n_res_blocks: Annotated[int | None, typer.Option("--stage1-n-res-blocks")] = None,
+    stage1_dropout: Annotated[float | None, typer.Option("--stage1-dropout")] = None,
+    stage2_generator: Annotated[Mode | None, typer.Option("--stage2-generator")] = None,
+    stage2_hidden_dim: Annotated[int | None, typer.Option("--stage2-hidden-dim")] = None,
+    stage2_n_res_blocks: Annotated[int | None, typer.Option("--stage2-n-res-blocks")] = None,
+    stage2_dropout: Annotated[float | None, typer.Option("--stage2-dropout")] = None,
+    stage2_decoder: Annotated[Decoder | None, typer.Option("--stage2-decoder")] = None,
+    stage2_k_max: Annotated[int | None, typer.Option("--stage2-k-max")] = None,
+    stage2_context_dim: Annotated[int | None, typer.Option("--stage2-context-dim")] = None,
+    stage2_stage1_context: Annotated[Stage1Context | None, typer.Option("--stage2-stage1-context")] = None,
+    stage1_init_from: Annotated[Path | None, typer.Option("--stage1-init-from")] = None,
+    stage1_freeze: Annotated[bool | None, typer.Option("--stage1-freeze/--no-stage1-freeze")] = None,
+    stage2_init_from: Annotated[Path | None, typer.Option("--stage2-init-from")] = None,
+    stage2_freeze: Annotated[bool | None, typer.Option("--stage2-freeze/--no-stage2-freeze")] = None,
+    conditioning: Annotated[Conditioning | None, typer.Option("--conditioning")] = None,
+    router: Annotated[bool | None, typer.Option("--router/--no-router")] = None,
+    router_type: Annotated[str | None, typer.Option("--router-type")] = None,
+    n_experts: Annotated[int | None, typer.Option("--n-experts")] = None,
+    router_axis: Annotated[list[str] | None, typer.Option("--router-axis")] = None,
    out: Annotated[
-        Optional[Path],
+        Path | None,
        typer.Option("--out", "-o", help="Run dir (default: auto from hyperparams)"),
    ] = None,
    comment: Annotated[
-        Optional[str],
+        str | None,
        typer.Option("--comment", help="Free-text note recorded in config.toml's meta section"),
    ] = None,
    data: Annotated[
-        Optional[Path],
+        Path | None,
        typer.Option(
            "--data",
            help="Dataset path to fill in the printed next-step command (not stored in the config)",
@@ -931,7 +930,7 @@ def new_run(
        # its stage1_model/stage2_model/conditioning content.
        "config_version": gconfig.CONFIG_VERSION,
        "git_hash": gconfig.git_hash(),
-        "created_at": datetime.now(timezone.utc).isoformat(timespec="seconds"),
+        "created_at": datetime.now(UTC).isoformat(timespec="seconds"),
        "created_by": "giant new-run",
    }
    if comment:
@@ -958,7 +957,7 @@ app.add_typer(model_app, name="model")
@model_app.command("summary")
 def model_summary(
    config: Annotated[
-        Optional[Path],
+        Path | None,
        typer.Option("--config", "-c", help="TOML config file (default: built-in defaults)"),
    ] = None,
    pdg_vocab: Annotated[
@@ -1017,7 +1016,7 @@ def predict(
        ),
    ] = Coord.global_,
    out: Annotated[
-        Optional[Path],
+        Path | None,
        typer.Option(
            "--out",
            "-o",
@@ -1050,11 +1049,11 @@ def predict(
        ),
    ] = Weights.raw,
    device: Annotated[
-        Optional[str],
+        str | None,
        typer.Option("--device", "-d", help="cpu | cuda | mps (default: auto)"),
    ] = None,
    comment: Annotated[
-        Optional[str],
+        str | None,
        typer.Option(
            "--comment",
            "-m",
@@ -1062,7 +1061,7 @@ def predict(
        ),
    ] = None,
    set_: Annotated[
-        Optional[list[str]],
+        list[str] | None,
        typer.Option(
            "--set",
            help="Override a sampling-only model_config key on this checkpoint, "
@@ -1092,7 +1091,7 @@ def predict(
    from giant.sample import resolve_n_sec, sample_stage1, sample_stage2

    batch_size_auto = False
-    batch_size_value: Optional[int] = None
+    batch_size_value: int | None = None
    if batch_size.strip().lower() == "auto":
        batch_size_auto = True
    else:
@@ -1272,7 +1271,7 @@ def predict(
            # "no snapping at inference"). "onehot"/"embedding": PDG
            # resolution IS the secondary's identity — see
            # decode_secondary_identity's docstring.
-            sec_E, sec_dir_world, sec_mass, sec_charge, sec_pdg_code, _l1_dist = decode_secondary_identity(
+            sec_E, sec_dir_world, _, _, sec_pdg_code, _l1_dist = decode_secondary_identity(
                sec_decoder,
                sec_cont,
                sec_type,
@@ -1463,28 +1462,28 @@ def rollout(
    ] = Weights.raw,
    batch_size: Annotated[int, typer.Option("--batch-size", "-b", help="Tracks stepped per model forward")] = 4096,
    max_tracks_per_event: Annotated[
-        Optional[int],
+        int | None,
        typer.Option(
            "--max-tracks-per-event",
            help="Safety cap on tracks per shower (sub-cap secondaries deposit in place)",
        ),
    ] = None,
    escape_threshold: Annotated[
-        Optional[float],
+        float | None,
        typer.Option(
            "--escape-threshold",
            help="Override the oracle's NN-distance escape threshold [mm]",
        ),
    ] = None,
-    n_events: Annotated[Optional[int], typer.Option("--n-events", help="Cap number of seed events")] = None,
-    device: Annotated[Optional[str], typer.Option("--device", "-d", help="cpu | cuda | mps (auto)")] = None,
-    out: Annotated[Optional[Path], typer.Option("--out", "-o", help="Output steps parquet")] = None,
+    n_events: Annotated[int | None, typer.Option("--n-events", help="Cap number of seed events")] = None,
+    device: Annotated[str | None, typer.Option("--device", "-d", help="cpu | cuda | mps (auto)")] = None,
+    out: Annotated[Path | None, typer.Option("--out", "-o", help="Output steps parquet")] = None,
    seed: Annotated[
-        Optional[int],
+        int | None,
        typer.Option("--seed", help="Torch/numpy seed for reproducibility"),
    ] = None,
    set_: Annotated[
-        Optional[list[str]],
+        list[str] | None,
        typer.Option(
            "--set",
            help="Override a sampling-only model_config key on this checkpoint, "
@@ -1505,7 +1504,8 @@ def rollout(
    from giant.checkpoint_io import CheckpointCompatibilityError, load_for_inference
    from giant.data.loader import find_parquet_files
    from giant.geometry import GeometryOracle
-    from giant.rollout import L1DistCollector, RolloutSummary, rollout as run_rollout
+    from giant.rollout import L1DistCollector, RolloutSummary
+    from giant.rollout import rollout as run_rollout

    _t_setup_start = time.perf_counter()

@@ -1640,7 +1640,7 @@ def rollout(
            "max_tracks_per_event": max_tracks_per_event,
            "escape_threshold": escape_threshold,
            "n_events": n_events,
-            "n_seed_events": int(len(seeds["event_id"])),
+            "n_seed_events": len(seeds["event_id"]),
            "weights": weights.value,
            "batch_size": batch_size,
            "device": str(_device),
@@ -1700,7 +1700,7 @@ def analyze_prep(
        ),
    ],
    label: Annotated[
-        Optional[list[str]],
+        list[str] | None,
        typer.Option(
            "--label",
            help="Series name for a rollout YAML, positionally matched to it — give none, "
@@ -1709,7 +1709,7 @@ def analyze_prep(
        ),
    ] = None,
    run_dir: Annotated[
-        Optional[Path],
+        Path | None,
        typer.Option(
            "--run-dir",
            "-o",
@@ -1797,7 +1797,7 @@ def analyze_render(
 def analyze_metrics(
    run_dir: Annotated[Path, typer.Argument(help="Run directory containing metrics.csv (from `giant train`)")],
    out_dir: Annotated[
-        Optional[Path],
+        Path | None,
        typer.Option(
            "--out",
            "-o",
@@ -1823,7 +1823,7 @@ def analyze_submit(
    ],
    accounting_group: Annotated[str, typer.Option("--accounting-group")],
    label: Annotated[
-        Optional[list[str]],
+        list[str] | None,
        typer.Option(
            "--label",
            help="Series name for a rollout YAML, positionally matched to it — give none, "
@@ -1832,7 +1832,7 @@ def analyze_submit(
        ),
    ] = None,
    run_dir: Annotated[
-        Optional[Path],
+        Path | None,
        typer.Option(
            "--run-dir",
            "-o",
@@ -9,7 +9,7 @@ import subprocess
 import sys
 import tomllib
 from dataclasses import dataclass, field
-from datetime import datetime, timezone
+from datetime import UTC, datetime
 from enum import Enum
 from pathlib import Path
 from typing import TYPE_CHECKING
@@ -80,7 +80,7 @@ class ConditioningAxisConfig:
    n_layers: int = 1

    @classmethod
-    def from_dict(cls, d: dict | None) -> "ConditioningAxisConfig":
+    def from_dict(cls, d: dict | None) -> ConditioningAxisConfig:
        d = d or {}
        return cls(
            type=d.get("type", "physical"),
@@ -106,7 +106,7 @@ class ConditioningConfig:
    material: ConditioningAxisConfig = field(default_factory=ConditioningAxisConfig)

    @classmethod
-    def from_dict(cls, d: dict | None) -> "ConditioningConfig":
+    def from_dict(cls, d: dict | None) -> ConditioningConfig:
        d = d or {}
        return cls(
            out_dim=d.get("out_dim", 128),
@@ -130,7 +130,7 @@ class FlowConfig:
    time_dim: int = 64

    @classmethod
-    def from_dict(cls, d: dict | None) -> "FlowConfig":
+    def from_dict(cls, d: dict | None) -> FlowConfig:
        d = d or {}
        return cls(time_dim=d.get("time_dim", 64))

@@ -144,7 +144,7 @@ class DdpmConfig:
    n_steps: int = 1000

    @classmethod
-    def from_dict(cls, d: dict | None) -> "DdpmConfig":
+    def from_dict(cls, d: dict | None) -> DdpmConfig:
        d = d or {}
        return cls(time_dim=d.get("time_dim", 64), n_steps=d.get("n_steps", 1000))

@@ -166,7 +166,7 @@ class Stage1WganConfig:
    critic_n_res_blocks: int = 0

    @classmethod
-    def from_dict(cls, d: dict | None) -> "Stage1WganConfig":
+    def from_dict(cls, d: dict | None) -> Stage1WganConfig:
        d = d or {}
        return cls(
            noise_dim=d.get("noise_dim", 64),
@@ -198,7 +198,7 @@ class Stage2WganConfig(Stage1WganConfig):
    gumbel_tau_end: float = 0.1

    @classmethod
-    def from_dict(cls, d: dict | None) -> "Stage2WganConfig":
+    def from_dict(cls, d: dict | None) -> Stage2WganConfig:
        d = d or {}
        return cls(
            noise_dim=d.get("noise_dim", 64),
@@ -293,7 +293,7 @@ class RouterConfig:
    extra: dict = field(default_factory=dict)

    @classmethod
-    def from_dict(cls, d: dict | None) -> "RouterConfig":
+    def from_dict(cls, d: dict | None) -> RouterConfig:
        d = d or {}
        return cls(
            enabled=d.get("enabled", False),
@@ -360,7 +360,7 @@ class TrunkConfig:
    block_conditioning: str = "add"

    @classmethod
-    def from_dict(cls, d: dict | None) -> "TrunkConfig":
+    def from_dict(cls, d: dict | None) -> TrunkConfig:
        d = d or {}
        return cls(type=d.get("type", "resmlp"), block_conditioning=d.get("block_conditioning", "add"))

@@ -377,7 +377,7 @@ class Stage2RouterConfig(RouterConfig):
    tie_to_stage1: bool = False

    @classmethod
-    def from_dict(cls, d: dict | None) -> "Stage2RouterConfig":
+    def from_dict(cls, d: dict | None) -> Stage2RouterConfig:
        d = d or {}
        known = _ROUTER_KNOWN_KEYS | {"tie_to_stage1"}
        return cls(
@@ -447,7 +447,7 @@ class NSecConfig:
    sampling: str = "greedy"

    @classmethod
-    def from_dict(cls, d: dict | None) -> "NSecConfig":
+    def from_dict(cls, d: dict | None) -> NSecConfig:
        d = d or {}
        return cls(
            mode=d.get("mode", "head"),
@@ -498,7 +498,7 @@ class ParticleTypeConfig:
    class_weighting: str = "none"

    @classmethod
-    def from_dict(cls, d: dict | None) -> "ParticleTypeConfig":
+    def from_dict(cls, d: dict | None) -> ParticleTypeConfig:
        d = d or {}
        return cls(
            target=d.get("target", "onehot"),
@@ -537,7 +537,7 @@ class AutoregressiveConfig:
    attn_n_layers: int = 2

    @classmethod
-    def from_dict(cls, d: dict | None) -> "AutoregressiveConfig":
+    def from_dict(cls, d: dict | None) -> AutoregressiveConfig:
        d = d or {}
        return cls(
            order=d.get("order", "energy_desc"),
@@ -576,7 +576,7 @@ class HeadConfig:
    depth: int = 2  # matches build_mlp_head's depth

    @classmethod
-    def from_dict(cls, d: dict | None) -> "HeadConfig":
+    def from_dict(cls, d: dict | None) -> HeadConfig:
        d = d or {}
        return cls(hidden_ratio=d.get("hidden_ratio", 0.5), depth=d.get("depth", 2))

@@ -593,7 +593,7 @@ class Stage1HeadsConfig:
    n_sec: HeadConfig = field(default_factory=HeadConfig)

    @classmethod
-    def from_dict(cls, d: dict | None) -> "Stage1HeadsConfig":
+    def from_dict(cls, d: dict | None) -> Stage1HeadsConfig:
        d = d or {}
        return cls(n_sec=HeadConfig.from_dict(d.get("n_sec")))

@@ -611,7 +611,7 @@ class Stage2HeadsConfig:
    type: HeadConfig = field(default_factory=HeadConfig)

    @classmethod
-    def from_dict(cls, d: dict | None) -> "Stage2HeadsConfig":
+    def from_dict(cls, d: dict | None) -> Stage2HeadsConfig:
        d = d or {}
        return cls(n_sec=HeadConfig.from_dict(d.get("n_sec")), type=HeadConfig.from_dict(d.get("type")))

@@ -659,7 +659,7 @@ class Stage1ModelConfig:
    heads: Stage1HeadsConfig = field(default_factory=Stage1HeadsConfig)

    @classmethod
-    def from_dict(cls, d: dict | None) -> "Stage1ModelConfig":
+    def from_dict(cls, d: dict | None) -> Stage1ModelConfig:
        d = d or {}
        return cls(
            active=d.get("active", True),
@@ -745,7 +745,7 @@ class Stage2ModelConfig:
    heads: Stage2HeadsConfig = field(default_factory=Stage2HeadsConfig)

    @classmethod
-    def from_dict(cls, d: dict | None) -> "Stage2ModelConfig":
+    def from_dict(cls, d: dict | None) -> Stage2ModelConfig:
        d = d or {}
        return cls(
            active=d.get("active", True),
@@ -839,7 +839,7 @@ class TrainConfig:
    precision: str = "fp32"

    @classmethod
-    def from_dict(cls, d: dict | None) -> "TrainConfig":
+    def from_dict(cls, d: dict | None) -> TrainConfig:
        d = d or {}
        return cls(
            epochs=d.get("epochs", 100),
@@ -895,7 +895,7 @@ class GiantConfig:
    train: TrainConfig = field(default_factory=TrainConfig)

    @classmethod
-    def from_dict(cls, d: dict | None) -> "GiantConfig":
+    def from_dict(cls, d: dict | None) -> GiantConfig:
        d = d or {}
        return cls(
            conditioning=ConditioningConfig.from_dict(d.get("conditioning")),
@@ -938,7 +938,7 @@ def leaf_paths(node: dict, prefix: str = "") -> list[str]:
 def git_hash() -> str:
    try:
        return subprocess.check_output(["git", "rev-parse", "HEAD"], stderr=subprocess.DEVNULL).decode().strip()
-    except Exception:
+    except Exception:  # noqa: BLE001 - any failure (no git, no repo, ...) degrades to "unknown"
        return "unknown"


@@ -1852,7 +1852,7 @@ def default_out_dir_name(cfg: dict, now: datetime | None = None) -> str:
    name unboundedly. This name doubles as the run's W&B id (see
    giant.training), which is the reason a timestamp is always included.
    """
-    now = now or datetime.now()
+    now = now or datetime.now()  # noqa: DTZ005 - human-readable local wall-clock time for run/W&B naming, not stored
    tokens = []
    overflow = []
    for label, candidate in _OUT_DIR_NAME_CANDIDATES:
@@ -1951,7 +1951,7 @@ def build_run_meta(
        "config_version": CONFIG_VERSION,
        "git_hash": git_hash(),
        "seed": seed,
-        "timestamp_utc": datetime.now(timezone.utc).isoformat(timespec="seconds"),
+        "timestamp_utc": datetime.now(UTC).isoformat(timespec="seconds"),
        "python_version": sys.version.split()[0],
        "torch_version": torch.__version__,
        "command": " ".join(sys.argv),
@@ -1,6 +1,7 @@
+from collections.abc import Iterator, Mapping
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import TYPE_CHECKING, Any, Iterator, Mapping
+from typing import TYPE_CHECKING, Any

 import numpy as np
 import polars as pl
@@ -172,7 +172,7 @@ class NormalizerEntry:
        }

    @classmethod
-    def from_json(cls, d: dict) -> "NormalizerEntry":
+    def from_json(cls, d: dict) -> NormalizerEntry:
        return cls(
            cond_norm=Normalizer.from_dict(d["cond_norm"]),
            tgt_norm=Normalizer.from_dict(d["tgt_norm"]),
@@ -194,7 +194,7 @@ class SetupCache:
    """Keyed by `topn_key(axis, n_classes)`."""

    @classmethod
-    def empty(cls, files: list[Path]) -> "SetupCache":
+    def empty(cls, files: list[Path]) -> SetupCache:
        return cls(fingerprint=fingerprint_files(files))

    def to_json(self) -> dict:
@@ -222,7 +222,7 @@ class SetupCache:
        return d

    @classmethod
-    def from_json(cls, d: dict) -> "SetupCache":
+    def from_json(cls, d: dict) -> SetupCache:
        vocab = None
        if "vocab" in d:
            pdg_map = {int(k): v for k, v in d["vocab"]["pdg_map"].items()}
@@ -247,7 +247,7 @@ class SetupCache:
            topn_maps=topn_maps,
        )

-    def merge(self, other: "SetupCache") -> "SetupCache":
+    def merge(self, other: SetupCache) -> SetupCache:
        """Union of both caches; `other`'s populated fields win on a shared key.

        Used by `save` to combine freshly-computed sections with whatever a
@@ -23,9 +23,10 @@ install stays lean.

 from __future__ import annotations

+from collections.abc import Iterable
 from dataclasses import dataclass, field
 from pathlib import Path
-from typing import Any, Iterable
+from typing import Any

 import numpy as np
 import polars as pl
@@ -155,7 +156,7 @@ class GeometryOracle:
        )

    @classmethod
-    def load(cls, path: str | Path) -> "GeometryOracle":
+    def load(cls, path: str | Path) -> GeometryOracle:
        _require_sklearn()
        import joblib

@@ -353,7 +354,7 @@ def _fit_slab_lookup(
        radius_max=radius_max,
    )
    info = {
-        "n_segments": int(len(materials)),
+        "n_segments": len(materials),
        "z_range": (z_min, z_max),
        "median_z_spacing": median_spacing,
        "radius_max": radius_max,
@@ -412,7 +413,7 @@ def build_geometry_oracle(
                "method": "slab",
                "depth_axis": depth_axis,
                "n_bins": n_bins,
-                "n_reference_points": int(len(pos)),
+                "n_reference_points": len(pos),
                "escape_factor": escape_factor,
                "n_files": len(files),
                **info,
@@ -458,7 +459,7 @@ def build_geometry_oracle(
        metadata={
            "method": method,
            "k": k,
-            "n_reference_points": int(len(X)),
+            "n_reference_points": len(X),
            "median_nn_dist": median_nn,
            "escape_factor": escape_factor,
            "n_files": len(files),
@@ -1,7 +1,7 @@
 """Factories: `build_models`/`build_critics` assemble the top-level stage
 models from a config dict (issues.md Issue 8)."""

-import torch.nn as nn
+from torch import nn

 from giant.config import ConditioningConfig, Stage1ModelConfig, Stage2ModelConfig
 from giant.constants import X_DIM
@@ -2,8 +2,8 @@
 identity (issues.md Issue 8)."""

 import torch
-import torch.nn as nn
 import torch.nn.functional as F
+from torch import nn

 from giant.cond_layout import CondLayout
 from giant.config import ConditioningAxisConfig
@@ -4,9 +4,10 @@ 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
+from torch import nn


 class HistoryEncoder(nn.Module):
@@ -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]] = {}

@@ -185,7 +197,7 @@ class AttentionHistory(HistoryEncoder):
        return self.in_proj(x)

    def forward(self, feat: torch.Tensor, has_prev: torch.Tensor) -> torch.Tensor:
-        B, K, _ = feat.shape
+        _, K, _ = feat.shape
        x = self._embed(feat, has_prev)
        mask = nn.Transformer.generate_square_subsequent_mask(K, device=feat.device)
        for block in self.blocks:
@@ -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]
@@ -4,7 +4,7 @@ no dependency on any other `giant.model` submodule (issues.md Issue 8)."""
 import math

 import torch
-import torch.nn as nn
+from torch import nn


 class SinusoidalEmbedding(nn.Module):
@@ -2,7 +2,7 @@
 `CriticModel` — composed from encoders/trunks/history (issues.md Issue 8)."""

 import torch
-import torch.nn as nn
+from torch import nn

 from giant.config import ConditioningAxisConfig, HeadConfig, ParticleTypeConfig
 from giant.constants import CONT_SLOT_DIM, K_MAX, PARTICLE_PHYS_DIM, SEC_DIM, SEC_SLOT_DIM, X_DIM
@@ -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,
@@ -79,9 +79,13 @@ from giant.model.trunks import (
 )

 __all__ = [
+    "BLOCK_REGISTRY",
+    "HISTORY_REGISTRY",
+    "OBJECTIVE_REGISTRY",
+    "ROUTER_REGISTRY",
+    "TRUNK_REGISTRY",
    "AdaLNResBlock",
    "AttentionHistory",
-    "BLOCK_REGISTRY",
    "ComposedRouter",
    "ConditionEncoder",
    "ContextAdapter",
@@ -91,17 +95,14 @@ __all__ = [
    "ExpertTrunk",
    "FilmResBlock",
    "FlowObjective",
-    "HISTORY_REGISTRY",
    "HistoryEncoder",
    "LinearTrunk",
    "MarkovHistory",
    "NoHistory",
    "NoneRouter",
-    "OBJECTIVE_REGISTRY",
    "Objective",
    "PdgRouter",
    "ProcessRouter",
-    "ROUTER_REGISTRY",
    "ResBlock",
    "RoutedTrunk",
    "Router",
@@ -110,7 +111,6 @@ __all__ = [
    "Stage2Autoregressive",
    "Stage2OneShot",
    "StageModel",
-    "TRUNK_REGISTRY",
    "Trunk",
    "WganObjective",
    "_CausalAttnBlock",
@@ -8,8 +8,8 @@ import re
 from collections.abc import Sequence

 import torch
-import torch.nn as nn
 import torch.nn.functional as F
+from torch import nn

 from giant.cond_layout import CondLayout
 from giant.constants import COND_DIM
@@ -398,7 +398,7 @@ def _build_router_from_cfg(
    """Resolve one stage's `router` config into a `Router`, single-axis or
    composed. `gumbel` is set as a post-construction attribute (shared by
    every router type, not a per-type constructor kwarg)."""
-    shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
+    shared_vocab = {"pdg_vocab": pdg_vocab, "mat_vocab": mat_vocab}
    if router_cfg["type"] == "composed":
        axes = _parse_composed_axes(router_cfg)
        _check_router_conditioning_compat([a["type"] for a in axes], particle_conditioning)
@@ -14,7 +14,7 @@ class CosineSchedule:
        betas = np.clip(1.0 - alpha_bars[1:] / alpha_bars[:-1], 0.0, 0.999).astype(np.float32)

        self.betas = torch.from_numpy(betas)
-        self.alphas = torch.from_numpy((1.0 - betas))
+        self.alphas = torch.from_numpy(1.0 - betas)
        self.alpha_bars = torch.from_numpy(alpha_bars[1:])

    def to(self, device: torch.device) -> "CosineSchedule":
@@ -35,7 +35,7 @@ finding those two tests independently converge on.
 import copy
 from dataclasses import dataclass, field

-import torch.nn as nn
+from torch import nn

 from giant.config import INFERENCE_OVERRIDES, _get_path, _set_path, leaf_paths
 from giant.model.builders import build_critics, build_models
@@ -224,7 +224,7 @@ def summarize_model(cfg: dict, pdg_vocab: int, mat_vocab: int) -> ModelSummary:
            _set_path(probe_cfg, path, candidate)
            try:
                changed = _fingerprint(_built_modules(probe_cfg, pdg_vocab, mat_vocab)) != baseline_fp
-            except Exception:
+            except Exception:  # noqa: BLE001 - a perturbation that fails to even build counts as "consumed"
                changed = True
            if changed:
                break
@@ -10,7 +10,7 @@ free — no separate "routed transformer trunk" class needed.
 """

 import torch
-import torch.nn as nn
+from torch import nn

 from giant.model.layers import build_block
 from giant.model.routers import Router
@@ -1,4 +1,4 @@
-from typing import Callable
+from collections.abc import Callable

 import torch

@@ -19,7 +19,7 @@ free-running history representation stays unsnapped — see

 from __future__ import annotations

-from functools import lru_cache
+from functools import cache
 from typing import TYPE_CHECKING

 import numpy as np
@@ -42,7 +42,7 @@ _CHARGE_WEIGHT = 50.0
 _LOG_EPS = 1e-8


-@lru_cache(maxsize=None)
+@cache
 def particle_mass_charge(pdg: int) -> tuple[float, float]:
    """Return (mass_MeV, charge_e) for a raw PDG code.

@@ -135,7 +135,7 @@ def invert_dense_map(m: dict[int, int]) -> dict[int, int]:

 def decode_topn_class(
    class_idx: np.ndarray,
-    topn_map: "TopNMap",
+    topn_map: TopNMap,
    n_classes: int,
    other_policy: str = "sample",
    rng: np.random.Generator | None = None,
@@ -13,6 +13,7 @@ from giant.constants import (
    X_DIM,
 )
 from giant.data import setup_cache
+from giant.data.dataset import StreamingStepsDataset, make_event_split
 from giant.data.loader import (
    TopNMap,
    _topn_plus_other_map,
@@ -23,13 +24,12 @@ from giant.data.loader import (
 from giant.data.scan import MetadataScan, ScanRequest, scan_metadata
 from giant.data.transforms import (
    Normalizer,
-    build_features,
-    _WelfordAccumulator,
    _ReservoirSampler,
+    _WelfordAccumulator,
+    build_features,
    sorted_membership,
 )
-from giant.data.dataset import make_event_split, StreamingStepsDataset
-from giant.model.network import build_models, build_critics, resolve_type_n_classes
+from giant.model.network import build_critics, build_models, resolve_type_n_classes
 from giant.training import train as run_training


@@ -17,7 +17,8 @@ treated as detector leakage and not deposited.
 from __future__ import annotations

 from collections import Counter
-from typing import TYPE_CHECKING, Callable, TypedDict
+from collections.abc import Callable
+from typing import TYPE_CHECKING, TypedDict

 import numpy as np
 import torch
@@ -117,7 +118,7 @@ def decode_secondary_identity(
    pre_dir: np.ndarray,
    sec_phys_norm: Normalizer,
    pdg_map: dict[int, int],
-    sec_type_topn_map: "TopNMap | None",
+    sec_type_topn_map: TopNMap | None,
    other_policy: str,
    rng: np.random.Generator | None,
 ) -> tuple[np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray, np.ndarray | None]:
@@ -390,34 +391,34 @@ def _terminal_rows(tr: dict[str, np.ndarray], sel: np.ndarray, reason: str, edep
    pos = tr["pre_pos"][sel]
    dir_ = tr["pre_dir"][sel]
    n = int(sel.sum())
-    return dict(
-        event_id=tr["event_id"][sel],
-        track_id=tr["track_id"][sel],
-        parent_id=tr["parent_id"][sel],
-        generation=tr["generation"][sel],
-        step_no=tr["step_in_track"][sel],
-        pdg=tr["pdg"][sel],
-        pre_x=pos[:, 0],
-        pre_y=pos[:, 1],
-        pre_z=pos[:, 2],
-        pre_E=tr["pre_E"][sel],
-        pre_dx=dir_[:, 0],
-        pre_dy=dir_[:, 1],
-        pre_dz=dir_[:, 2],
-        post_x=pos[:, 0],
-        post_y=pos[:, 1],
-        post_z=pos[:, 2],
-        post_E=np.zeros(n),
-        post_dx=dir_[:, 0],
-        post_dy=dir_[:, 1],
-        post_dz=dir_[:, 2],
-        edep=np.asarray(edep, dtype=np.float64).reshape(n),
-        step_length=np.zeros(n),
-        material=tr.get("_material", np.full(len(sel), "", dtype=object))[sel],
-        layer_id=tr.get("_layer_id", np.zeros(len(sel), dtype=np.int64))[sel],
-        n_sec_pred=np.zeros(n, dtype=np.int64),
-        termination_reason=np.full(n, reason, dtype=object),
-    )
+    return {
+        "event_id": tr["event_id"][sel],
+        "track_id": tr["track_id"][sel],
+        "parent_id": tr["parent_id"][sel],
+        "generation": tr["generation"][sel],
+        "step_no": tr["step_in_track"][sel],
+        "pdg": tr["pdg"][sel],
+        "pre_x": pos[:, 0],
+        "pre_y": pos[:, 1],
+        "pre_z": pos[:, 2],
+        "pre_E": tr["pre_E"][sel],
+        "pre_dx": dir_[:, 0],
+        "pre_dy": dir_[:, 1],
+        "pre_dz": dir_[:, 2],
+        "post_x": pos[:, 0],
+        "post_y": pos[:, 1],
+        "post_z": pos[:, 2],
+        "post_E": np.zeros(n),
+        "post_dx": dir_[:, 0],
+        "post_dy": dir_[:, 1],
+        "post_dz": dir_[:, 2],
+        "edep": np.asarray(edep, dtype=np.float64).reshape(n),
+        "step_length": np.zeros(n),
+        "material": tr.get("_material", np.full(len(sel), "", dtype=object))[sel],
+        "layer_id": tr.get("_layer_id", np.zeros(len(sel), dtype=np.int64))[sel],
+        "n_sec_pred": np.zeros(n, dtype=np.int64),
+        "termination_reason": np.full(n, reason, dtype=object),
+    }


@torch.no_grad()
@@ -442,14 +443,14 @@ def rollout(
    on_chunk: Callable[[dict[str, np.ndarray]], None] | None = None,
    particle_conditioning: str = "embedding",
    material_conditioning: str = "embedding",
-    pdg_topn_map: "TopNMap | None" = None,
-    mat_topn_map: "TopNMap | None" = None,
-    sec_type_topn_map: "TopNMap | None" = None,
+    pdg_topn_map: TopNMap | None = None,
+    mat_topn_map: TopNMap | None = None,
+    sec_type_topn_map: TopNMap | None = None,
    other_policy: str = "sample",
    seed: int | None = None,
    stage1_ddpm_steps: int = 1000,
    stage2_ddpm_steps: int = 1000,
-    l1_dist_collector: "L1DistCollector | None" = None,
+    l1_dist_collector: L1DistCollector | None = None,
 ) -> dict[str, np.ndarray] | RolloutSummary:
    """Run showers to completion.

@@ -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
-    once, so a flow/ddpm AR run costs ~`k_max * steps` model calls per
-    physics step (or ~`n_sec * steps` under `n_sec_pred=None` below, once
-    every row in the batch has stopped).
+    the "K sequential forwards" cost applies per-token here, not once, so a
+    flow/ddpm AR run costs ~`n_sec * steps` model calls per physics step
+    (measured on `configs/baseline.toml`: 0.382 secondaries/step at rollout
+    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
-    resolved by `resolve_n_sec` — `n_sec.mode` in `("head", "truth")`, or a
-    stop-token decoder driven by `_assemble_stage2_ar_inputs_scheduled`'s
-    ground-truth `n_sec`, which must run the *full* `k_max`-length free-
-    running self-sample regardless of the decoder's own stop head — the
-    scheduled-sampling training contract does not truncate). This always
-    runs the full `k_max`-iteration loop, masking by the given count at the
-    end exactly as before.
+    `n_sec_pred`, if given (as resolved by `resolve_n_sec` — `n_sec.mode` in
+    `("head", "truth")`, or a stop-token decoder driven by
+    `_assemble_stage2_ar_inputs_scheduled`'s ground-truth `n_sec`, which must
+    run the *full* `k_max`-length free-running self-sample regardless of the
+    decoder's own stop head — the scheduled-sampling training contract does
+    not truncate, see `full_length` below) fixes each row's secondary count
+    up front.

    `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
-    the batch has fired, the loop breaks before spending a model call on the
-    next slot's token — the average-case cost win the docstring above
-    describes. A row that never fires within `k_max` is capped there
-    (`K_MAX` stays a safety cap, not a modeling ceiling).
+    `n_sec_pred` is the first slot index where this fires. 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)
-        history_feat = prev_repr.unsqueeze(1)  # (B, 1, CONT_SLOT_DIM + type_dim)
-        remaining_frac = remaining.unsqueeze(1)  # (B, 1)
-        slot_idx = torch.full((B, 1), k / max(k_max - 1, 1), device=device, dtype=torch.float32)
+        if active_idx.numel() == 0:
+            break
+        Bc = active_idx.numel()
+        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,
-                cond_cat,
-                stage1_out,
+                cc,
+                ck,
+                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
-            finished = finished | stop_now
-            if finished.all():
-                break
+            newly_stopped = stop_now & ~finished.index_select(0, active_idx)
+            derived_n_sec[active_idx[newly_stopped]] = k
+            finished[active_idx[stop_now]] = True
+            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,
-                cond_cat,
-                stage1_out,
+                cc,
+                ck,
+                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,
-                    cond_cat,
-                    stage1_out,
+                    cc,
+                    ck,
+                    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,
-                cond_cat,
-                stage1_out,
+                cc,
+                ck,
+                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_type[:, k] = type_k
+        sec_cont[active_idx, k] = cont_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
@@ -82,7 +82,7 @@ def _max_index(parent: Path, pattern: re.Pattern) -> int:

 def _git_user_name() -> str | None:
    try:
-        out = subprocess.run(["git", "config", "user.name"], capture_output=True, text=True, timeout=2)
+        out = subprocess.run(["git", "config", "user.name"], capture_output=True, text=True, timeout=2, check=False)
    except (OSError, subprocess.SubprocessError):
        # OSError (e.g. git not on PATH) and subprocess.SubprocessError
        # (e.g. TimeoutExpired) are unrelated hierarchies — TimeoutExpired
@@ -604,7 +604,7 @@ def _run_bump(
    if not root_path.is_dir():
        raise SystemExit(f"error: {root_path} is not a directory")

-    date = date or dt.date.today().isoformat()
+    date = date or dt.date.today().isoformat()  # noqa: DTZ011 - local calendar date for the dataset-version log, not stored
    by = by if by is not None else _git_user_name()
    if gen is None:
        new_dirs, log_line = plan_bump_gen(root_path, kind, reason, by, date, to)
@@ -31,9 +31,9 @@ import torch

 sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

-from giant.constants import COND_DIM, SEC_SLOT_DIM, X_DIM  # noqa: E402
-from giant.model import network as net  # noqa: E402
-from tests.legacy import network_v02_snapshot as legacy  # noqa: E402
+from giant.constants import COND_DIM, SEC_SLOT_DIM, X_DIM
+from giant.model import network as net
+from tests.legacy import network_v02_snapshot as legacy


 def _random_batch(model_config: dict, batch: int, seed: int):
@@ -28,8 +28,8 @@ import subprocess
 import sys
 import uuid
 import zlib
-from dataclasses import dataclass
 from concurrent.futures import ThreadPoolExecutor, as_completed
+from dataclasses import dataclass
 from pathlib import Path

 # Must match giant/tools/bump_dataset_version.py's GEN_RE.
@@ -147,7 +147,7 @@ def run_job(
    cmd = build_cmd(executable, job, events_per_file, energy_gev)

    env = dict(os.environ, MINICALOSIM_SEED=str(job_seed(kind, gen, job, energy_gev)))
-    result = subprocess.run(cmd, cwd=workdir, capture_output=True, text=True, env=env)
+    result = subprocess.run(cmd, cwd=workdir, capture_output=True, text=True, env=env, check=False)

    if result.returncode != 0:
        return JobResult(
@@ -10,10 +10,9 @@ from __future__ import annotations
 import os
 from enum import Enum
 from pathlib import Path
-from typing import Optional
+from typing import Annotated

 import typer
-from typing_extensions import Annotated

 from giant.config import Conditioning

@@ -73,7 +72,7 @@ class PoolType(str, Enum):
 def convert(
    root_files: Annotated[list[Path], typer.Argument(help="Input ROOT file(s)")],
    output: Annotated[
-        Optional[Path],
+        Path | None,
        typer.Option(
            "--output",
            "-o",
@@ -108,7 +107,7 @@ def convert(
        ),
    ] = _DATASET_ROOT_DEFAULT,
    schema: Annotated[
-        Optional[str],
+        str | None,
        typer.Option(
            "--schema",
            help="Schema tag to write parquets under, e.g. schema2 (only used "
@@ -192,13 +191,13 @@ def migrate(
 def bump_gen(
    reason: Annotated[str, typer.Option("--reason", help="Why this gen exists")],
    kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
-    by: Annotated[Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")] = None,
+    by: Annotated[str | None, typer.Option("--by", help="Attribution (default: git user.name)")] = None,
    date: Annotated[
-        Optional[str],
+        str | None,
        typer.Option("--date", help="Override date (default: today, ISO)"),
    ] = None,
    to: Annotated[
-        Optional[str],
+        str | None,
        typer.Option(
            "--to",
            metavar="genN",
@@ -227,13 +226,13 @@ def bump_schema(
    gen: Annotated[str, typer.Option("--gen", help="Existing gen tag, e.g. gen1")],
    reason: Annotated[str, typer.Option("--reason", help="Why this schema exists")],
    kind: Annotated[str, typer.Option("--kind", help="steps | hits | ... (default: steps)")] = "steps",
-    by: Annotated[Optional[str], typer.Option("--by", help="Attribution (default: git user.name)")] = None,
+    by: Annotated[str | None, typer.Option("--by", help="Attribution (default: git user.name)")] = None,
    date: Annotated[
-        Optional[str],
+        str | None,
        typer.Option("--date", help="Override date (default: today, ISO)"),
    ] = None,
    to: Annotated[
-        Optional[str],
+        str | None,
        typer.Option(
            "--to",
            metavar="schemaN",
@@ -272,7 +271,7 @@ def status(
 def update_manifest(
    manifests: Annotated[list[Path], typer.Argument(help="One or more .manifest files to update")],
    schema: Annotated[
-        Optional[str],
+        str | None,
        typer.Option(
            "--schema",
            metavar="schemaN",
@@ -280,7 +279,7 @@ def update_manifest(
        ),
    ] = None,
    gen: Annotated[
-        Optional[str],
+        str | None,
        typer.Option("--gen", metavar="genN", help="Target gen tag (default: keep existing gen)"),
    ] = None,
    execute: Annotated[
@@ -298,11 +297,11 @@ def update_manifest(
 def create_manifest(
    files: Annotated[list[Path], typer.Argument(help="Parquet files to include")],
    output: Annotated[
-        Optional[Path],
+        Path | None,
        typer.Option("--output", "-o", help="Explicit path for the new .manifest file"),
    ] = None,
    pool: Annotated[
-        Optional[str],
+        str | None,
        typer.Option(
            "--pool",
            metavar="DETECTOR",
@@ -310,7 +309,7 @@ def create_manifest(
        ),
    ] = None,
    type_: Annotated[
-        Optional[PoolType],
+        PoolType | None,
        typer.Option("--type", help="Pool type — full, holdout, or dev (required with --pool)"),
    ] = None,
    root: Annotated[Path, typer.Option("--root", help="Dataset root (used with --pool)")] = _DATASET_ROOT_DEFAULT,
@@ -459,7 +458,7 @@ def warm_cache(
        typer.Argument(help="Parquet file, directory, or .manifest — same as `giant train`'s"),
    ],
    config: Annotated[
-        Optional[Path],
+        Path | None,
        typer.Option(
            "--config",
            "-c",
@@ -469,7 +468,7 @@ def warm_cache(
        ),
    ] = None,
    val_fraction: Annotated[
-        Optional[float],
+        float | None,
        typer.Option(
            "--val-fraction",
            "-f",
@@ -477,7 +476,7 @@ def warm_cache(
        ),
    ] = None,
    seed: Annotated[
-        Optional[int],
+        int | None,
        typer.Option(
            "--seed",
            "-s",
@@ -485,7 +484,7 @@ def warm_cache(
        ),
    ] = None,
    particle_conditioning: Annotated[
-        Optional[Conditioning],
+        Conditioning | None,
        typer.Option(
            "--particle-conditioning",
            help="Must match the `giant train` run(s)' conditioning.particle.type to warm for. "
@@ -493,7 +492,7 @@ def warm_cache(
        ),
    ] = None,
    material_conditioning: Annotated[
-        Optional[Conditioning],
+        Conditioning | None,
        typer.Option(
            "--material-conditioning",
            help="Must match the `giant train` run(s)' conditioning.material.type "
@@ -502,7 +501,7 @@ def warm_cache(
        ),
    ] = None,
    router: Annotated[
-        Optional[bool],
+        bool | None,
        typer.Option(
            "--router/--no-router",
            help="Warm the process vocabulary too (only takes effect with --router-type process). "
@@ -510,11 +509,11 @@ def warm_cache(
        ),
    ] = None,
    router_type: Annotated[
-        Optional[str],
+        str | None,
        typer.Option("--router-type", help="Router implementation name. Not allowed together with --config"),
    ] = None,
    n_experts: Annotated[
-        Optional[int],
+        int | None,
        typer.Option("--n-experts", help="Number of routed experts. Not allowed together with --config"),
    ] = None,
    rebuild: Annotated[
@@ -144,7 +144,7 @@ def run_hparam_scan(
        start = time.monotonic()
        try:
            with open(out_dir / "train.log", "a") as log:
-                subprocess.run(cmd, env=env, stdout=log, stderr=subprocess.STDOUT)
+                subprocess.run(cmd, env=env, stdout=log, stderr=subprocess.STDOUT, check=False)
        except KeyboardInterrupt:
            print(
                f"\ninterrupted during {name} — re-run this script to resume "
@@ -32,7 +32,7 @@ MANIFEST_SUFFIX = ".manifest"
 # since today's pool assignment is encoded only by *which folder a file's
 # parquet was copied into* — not by anything in the filename itself.
 POOL_ASSIGNMENT: dict[str, dict[str, range | list[int]]] = {
-    "pbwo4": {"full": range(0, 6), "holdout": range(6, 10)},
+    "pbwo4": {"full": range(6), "holdout": range(6, 10)},
    "sampling_fe_scint": {"dev": [0], "full": [1, 2], "holdout": [3]},
    "sampling_pb_lar": {"dev": [0], "full": [1, 2], "holdout": [3]},
    "sampling_pb_scint": {"dev": [0], "full": [1, 2], "holdout": [3]},
@@ -106,7 +106,7 @@ def _convert_one(
    if output_path is not None:
        output_path.parent.mkdir(parents=True, exist_ok=True)
        cmd += ["--output", str(output_path)]
-    result = subprocess.run(cmd, capture_output=True, text=True)
+    result = subprocess.run(cmd, capture_output=True, text=True, check=False)
    return root_file, result.returncode, result.stdout, result.stderr


@@ -6,8 +6,8 @@ that tests and tooling construct directly.
 """

 from giant.training.checkpoint import build_checkpoint, init_stages_from_checkpoints, load_checkpoint
-from giant.training.metrics import MetricsCollector, MetricSpec
 from giant.training.loop import train
+from giant.training.metrics import MetricsCollector, MetricSpec
 from giant.training.trainers import (
    FlowDDPMStageTrainer,
    StageSpec,
@@ -9,9 +9,10 @@
 import os
 import signal
 import time
+from collections.abc import Callable
 from pathlib import Path
 from types import FrameType
-from typing import Callable
+from typing import Self

 import numpy as np
 import torch
@@ -47,7 +48,7 @@ class _GracefulShutdown:
            Callable[[int, FrameType | None], object] | signal.Handlers | int | None,
        ] = {}

-    def __enter__(self) -> "_GracefulShutdown":
+    def __enter__(self) -> Self:
        for sig in _CATCHABLE_SIGNALS:
            self._previous[sig] = signal.getsignal(sig)
            signal.signal(sig, self._handle)
@@ -188,7 +188,7 @@ class MetricsCollector:
        self.fieldnames = self._build_fieldnames()
        metrics_path = out_dir / "metrics.csv"
        append = resume and metrics_path.exists()
-        self._file = open(metrics_path, "a" if append else "w", newline="")
+        self._file = open(metrics_path, "a" if append else "w", newline="")  # noqa: SIM115 - kept open for the object's lifetime, closed in .close()
        self._writer = csv.DictWriter(self._file, fieldnames=self.fieldnames)
        if not append:
            self._writer.writeheader()
@@ -47,7 +47,7 @@ class MetricsTable:
    columns: dict[str, list[float]]

    @classmethod
-    def load(cls, path: str | Path) -> "MetricsTable":
+    def load(cls, path: str | Path) -> MetricsTable:
        with open(path, newline="") as f:
            rows = list(csv.DictReader(f))
        epochs = [int(float(r["epoch"])) for r in rows]
@@ -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
-    cheap one-step proxy, so building it costs the same `k_max` (`* steps`
-    for flow) sequential forwards `sample.py` pays at inference, EVERY batch
+    `giant.sample.sample_secondaries_ar` under `torch.no_grad()`, called here
+    with `full_length=True` — not a cheap one-step proxy, so building it
+    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()
@@ -20,7 +20,7 @@ from typing import NamedTuple

 import torch
 import torch.nn.functional as F
-import torch.optim as optim
+from torch import optim

 from giant.config import ParticleTypeConfig, Stage1ModelConfig, Stage2ModelConfig, TrainConfig
 from giant.constants import CONT_SLOT_DIM
@@ -333,7 +333,7 @@ class StageTrainer:
        #: "sampled"` — the stage-1 `StageTrainer` this (stage-2) trainer
        #: draws its context sample from. `None` for stage 1 itself, and for
        #: stage 2 under "truth".
-        self.stage1_source: "StageTrainer | None" = None
+        self.stage1_source: StageTrainer | None = None

    def attach_stage1(self, stage1_trainer: "StageTrainer") -> None:
        """Wires this (stage-2) trainer to the stage-1 trainer it should
@@ -1,13 +1,13 @@
 [project]
 name = "giant"
-version = "0.3.17"
+version = "0.3.20"
 description = "Geant4 step-function surrogate via conditional flow matching"
 readme = "README.md"
 requires-python = ">=3.12"
 dependencies = [
    "numpy>=1.26,<3",
    "polars>=1.0,<2",
-    "pyarrow>=16,<25",
+    "pyarrow>=16,<26",
    "tqdm>=4.60,<5",
    "typer>=0.12,<1",
    "pyyaml>=6,<7",
@@ -17,8 +17,8 @@ import re
 from collections.abc import Sequence

 import torch
-import torch.nn as nn
 import torch.nn.functional as F
+from torch import nn

 from giant.constants import (
    COND_DIM,
@@ -949,7 +949,7 @@ def _check_router_conditioning_compat(router_types: list[str], conditioning: str


 def _build_router_from_cfg(router_cfg: dict, pdg_vocab: int, mat_vocab: int, conditioning: str = "embedding") -> Router:
-    shared_vocab = dict(pdg_vocab=pdg_vocab, mat_vocab=mat_vocab)
+    shared_vocab = {"pdg_vocab": pdg_vocab, "mat_vocab": mat_vocab}
    if router_cfg["type"] == "composed":
        axes = _parse_composed_axes(router_cfg)
        _check_router_conditioning_compat([a["type"] for a in axes], conditioning)
@@ -977,15 +977,15 @@ def build_models(model_config: dict) -> tuple[nn.Module, nn.Module]:
    if router_cfg and router_cfg.get("enabled"):
        pdg_vocab = model_config["pdg_vocab"]
        mat_vocab = model_config["mat_vocab"]
-        shared = dict(
-            pdg_vocab=pdg_vocab,
-            mat_vocab=mat_vocab,
-            expert_hidden_dim=model_config.get("expert_hidden_dim") or model_config.get("hidden_dim", 128),
-            expert_n_blocks=model_config.get("expert_n_blocks") or model_config.get("n_blocks", 3),
-            emb_dim=model_config.get("emb_dim", EMB_DIM),
-            dropout=model_config.get("dropout", 0.1),
-            conditioning=model_config.get("conditioning", "embedding"),
-        )
+        shared = {
+            "pdg_vocab": pdg_vocab,
+            "mat_vocab": mat_vocab,
+            "expert_hidden_dim": model_config.get("expert_hidden_dim") or model_config.get("hidden_dim", 128),
+            "expert_n_blocks": model_config.get("expert_n_blocks") or model_config.get("n_blocks", 3),
+            "emb_dim": model_config.get("emb_dim", EMB_DIM),
+            "dropout": model_config.get("dropout", 0.1),
+            "conditioning": model_config.get("conditioning", "embedding"),
+        }
        conditioning = shared["conditioning"]
        stage1 = RoutedDenoisingMLP(
            router=_build_router_from_cfg(router_cfg, pdg_vocab, mat_vocab, conditioning),
@@ -5,12 +5,12 @@ from pathlib import Path

 import pytest
 import torch
+from test_train import _base_cfg, _run_train

 from giant.model.routers import EnergyRouter
 from giant.model.wgan import gradient_penalty
 from giant.training.amp import resolve_autocast
 from giant.training.stage2_inputs import _remaining_energy_fraction
-from test_train import _base_cfg, _run_train

 # ---------------------------------------------------------------------------
 # resolve_autocast
@@ -184,7 +184,7 @@ def test_weighted_profile_matches_manual_bincount():
    ea = R.entry_axis(lf)
    lf2 = R.attach_entry_axis(lf, ea)
    edges = np.linspace(0.0, 3.0, 4)  # depth bins along +z
-    mean, std = R.weighted_profile(lf2, R.depth_expr(), edges, pl.col("edep"))
+    mean, _ = R.weighted_profile(lf2, R.depth_expr(), edges, pl.col("edep"))
    assert mean.shape == (3,)
    # totals conserved: sum over bins == mean total edep per event
    assert np.isclose(mean.sum() * 1, (90.0 + 30.0) / 2)  # 2 events
@@ -197,7 +197,7 @@ def test_update_manifest_reports_missing_targets(tmp_path):
    # schema2 dir exists but the parquet file does not
    (tmp_path / "processed" / "steps" / "gen1" / "schema2").mkdir(parents=True)

-    lines, missing = plan_update_manifest(manifest, "schema2")
+    _, missing = plan_update_manifest(manifest, "schema2")
    assert len(missing) == 1
    assert "schema2" in str(missing[0])

@@ -313,7 +313,7 @@ def test_create_manifest_writes_relative_paths(tmp_path):
 def test_create_manifest_reports_missing_files(tmp_path):
    ghost = tmp_path / "processed" / "gen1" / "schema2" / "shard-000.parquet"
    output = tmp_path / "pools" / "full.manifest"
-    lines, missing, _ = plan_create_manifest(output, [ghost])
+    _, missing, _ = plan_create_manifest(output, [ghost])
    assert len(missing) == 1
    assert missing[0] == ghost.resolve()

@@ -9,7 +9,7 @@ from pathlib import Path

 from typer.testing import CliRunner

-import giant.cli as cli
+from giant import cli

 runner = CliRunner()

@@ -1,4 +1,5 @@
 import pytest
+
 from giant.cond_layout import AXIS_TYPES, CondLayout
 from giant.constants import COND_DIM, COND_DIM_BASE, MATERIAL_PHYS_DIM, PARTICLE_PHYS_DIM

@@ -576,7 +576,7 @@ def test_save_config_round_trips_three_level_nesting(tmp_path):
 # default_out_dir_name
 # ---------------------------------------------------------------------------

-_NOW = datetime(2026, 7, 29, 14, 30)
+_NOW = datetime(2026, 7, 29, 14, 30)  # noqa: DTZ001 - naive, matching default_out_dir_name's naive datetime.now()


 def _cfg_with(**dotted_overrides):
@@ -90,7 +90,7 @@ def _collect_names(source: str, filename: str) -> set[str]:
            names.add(node.attr)
        elif isinstance(node, ast.Constant) and isinstance(node.value, str) and id(node) not in docstring_ids:
            names.add(node.value)
-        elif isinstance(node, ast.arg):
+        elif isinstance(node, ast.arg):  # noqa: SIM114 - kept separate so ty narrows node.arg to str, not str | None
            names.add(node.arg)
        elif isinstance(node, ast.keyword) and node.arg is not None:
            names.add(node.arg)
@@ -1,3 +1,4 @@
+from test_pipeline import _make_synthetic_steps
 from typer.testing import CliRunner

 from giant import cli as giant_cli
@@ -5,7 +6,6 @@ from giant.config import Conditioning
 from giant.data import setup_cache
 from giant.tools import dwarf
 from giant.tools.dwarf import app
-from test_pipeline import _make_synthetic_steps

 runner = CliRunner()

@@ -1,9 +1,10 @@
 import torch
+
 from giant.config import ConditioningAxisConfig
 from giant.constants import COND_DIM
 from giant.model.network import Stage1Model
 from giant.model.schedule import CosineSchedule, flow_matching_loss
-from giant.sample import sample_flow, sample_ddim
+from giant.sample import sample_ddim, sample_flow

 PARTICLE_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
 MATERIAL_CFG = ConditioningAxisConfig(type="physical", emb_dim=8, n_layers=1)
@@ -2,8 +2,9 @@ import copy

 import pytest
 import torch
+
 from giant import config as gconfig
-from giant.constants import CONT_SLOT_DIM, COND_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM
+from giant.constants import COND_DIM, CONT_SLOT_DIM, PARTICLE_PHYS_DIM, SEC_SLOT_DIM
 from giant.model.network import (
    HISTORY_REGISTRY,
    AttentionHistory,
@@ -1217,18 +1218,18 @@ def test_stage_classes_are_stagemodel_subclasses(cls):
@pytest.mark.parametrize("cls", [Stage1Model, Stage2OneShot, Stage2Autoregressive])
@pytest.mark.parametrize("generator", ["flow", "ddpm", "wgan"])
 def test_stagemodel_time_emb_matches_objective_needs_time(cls, generator):
-    kwargs = dict(
-        pdg_vocab=5,
-        mat_vocab=3,
-        particle_cfg=PARTICLE_CFG,
-        material_cfg=MATERIAL_CFG,
-        hidden_dim=_STAGE_HIDDEN_DIM,
-        n_res_blocks=_STAGE_N_BLOCKS,
-        cond_out_dim=_STAGE_COND_OUT_DIM,
-        generator=generator,
-        time_dim=8,
-        noise_dim=8,
-    )
+    kwargs = {
+        "pdg_vocab": 5,
+        "mat_vocab": 3,
+        "particle_cfg": PARTICLE_CFG,
+        "material_cfg": MATERIAL_CFG,
+        "hidden_dim": _STAGE_HIDDEN_DIM,
+        "n_res_blocks": _STAGE_N_BLOCKS,
+        "cond_out_dim": _STAGE_COND_OUT_DIM,
+        "generator": generator,
+        "time_dim": 8,
+        "noise_dim": 8,
+    }
    if cls is Stage1Model:
        kwargs["n_sec_head_k_max"] = 15
    else:
@@ -20,7 +20,6 @@ from giant.model.schedule import (
 )
 from giant.sample import sample_secondaries

-
 # ── helpers ──────────────────────────────────────────────────────────────────


@@ -341,7 +340,7 @@ def test_encode_secondaries_energy_conservation():
 def test_encode_secondaries_stick_logits_match_naive_reference():
    """Cumsum-based remaining-budget computation must match a naive
    per-row, per-slot Python reference (no cumsum) within float tolerance."""
-    from giant.data.transforms import encode_secondaries, _EPS, _STICK_LOGIT_CLIP
+    from giant.data.transforms import _EPS, _STICK_LOGIT_CLIP, encode_secondaries

    rng = np.random.default_rng(11)
    N = 25
@@ -569,7 +568,7 @@ def test_decode_secondaries_degenerate_row_falls_back_to_even_split():
    e_sec = np.array([0.0, 4.0, 9.0, 30.0], dtype=np.float32)
    pre_dir = np.tile([0.0, 0.0, 1.0], (N, 1)).astype(np.float32)

-    sec_E, _sec_dir, _mass, _charge, sec_valid = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)
+    sec_E, _sec_dir, _mass, _charge, _ = decode_secondaries(sec_cont, n_sec, e_sec, pre_dir)

    for i, k in enumerate(n_sec):
        if k == 0:
@@ -593,7 +592,7 @@ def test_decode_secondaries_rescale_preserves_relative_shares():
    n_sec = np.array([4])
    pre_dir = np.array([[0.0, 0.0, 1.0]], dtype=np.float32)

-    sec_E_small, _, _, _, sec_valid = decode_secondaries(sec_cont, n_sec, np.array([5.0], dtype=np.float32), pre_dir)
+    sec_E_small, _, _, _, _ = decode_secondaries(sec_cont, n_sec, np.array([5.0], dtype=np.float32), pre_dir)
    sec_E_large, _, _, _, _ = decode_secondaries(sec_cont, n_sec, np.array([50.0], dtype=np.float32), pre_dir)

    ratio_small = sec_E_small[0, :4] / sec_E_small[0, 0]
@@ -9,8 +9,8 @@ import pytest

 pytest.importorskip("plotstyle")

-from giant.analysis import render as render_mod  # noqa: E402
-from giant.analysis.reduced import Reduced  # noqa: E402
+from giant.analysis import render as render_mod
+from giant.analysis.reduced import Reduced


 def _try_render(reduced: list[Reduced], out: Path):
@@ -9,7 +9,7 @@ import pytest
 import torch

 from giant.config import ConditioningAxisConfig, ParticleTypeConfig
-from giant.constants import TERM_ESCAPED, TERM_MAX_STEPS, TERM_UNKNOWN_PDG, K_MAX
+from giant.constants import K_MAX, TERM_ESCAPED, TERM_MAX_STEPS, TERM_UNKNOWN_PDG
 from giant.data.loader import TopNMap
 from giant.data.transforms import Normalizer
 from giant.model.network import (
@@ -21,7 +21,7 @@ from giant.model.network import (
 from giant.rollout import L1DistCollector, make_seed_frontier, rollout

 pytest.importorskip("sklearn")
-from giant import geometry as g  # noqa: E402
+from giant import geometry as g

 PDG_MAP = {22: 0, 11: 1, -11: 2}
 MAT_MAP = {"G4_AIR": 0, "G4_PbWO4": 1}
@@ -1,5 +1,7 @@
 """Tests for the mixture-of-experts routing prototype (giant/model/network.py)."""

+import itertools
+
 import pytest
 import torch

@@ -7,6 +9,7 @@ from giant.config import ConditioningAxisConfig
 from giant.constants import COND_DIM, K_MAX, SEC_DIM, X_DIM
 from giant.model.network import (
    BLOCK_REGISTRY,
+    ROUTER_REGISTRY,
    TRUNK_REGISTRY,
    AdaLNResBlock,
    ComposedRouter,
@@ -17,7 +20,6 @@ from giant.model.network import (
    NoneRouter,
    PdgRouter,
    ProcessRouter,
-    ROUTER_REGISTRY,
    ResBlock,
    RoutedTrunk,
    Stage1Model,
@@ -397,7 +399,7 @@ def test_energy_router_own_width_controls_own_coverage_independent_of_others():
            router.raw_width[0] = raw
            shares.append(router.gate(cond_cont, cond_cat)[0, 0].item())

-    assert all(a <= b + 1e-6 for a, b in zip(shares, shares[1:]))
+    assert all(a <= b + 1e-6 for a, b in itertools.pairwise(shares))


 def test_build_router_threads_learn_width_kwargs_through():
@@ -1010,9 +1012,7 @@ def test_build_models_routed_pair_composed_router_is_drop_in_for_sample_flow():
    n_sec_pred = stage2.predict_n_sec(cond_cont, cond_cat, stage1_norm).argmax(dim=-1)
    assert n_sec_pred.shape == (B,)

-    sec_cont, sec_type_emb, sec_valid = sample_secondaries(
-        stage2, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2
-    )
+    sec_cont, _, sec_valid = sample_secondaries(stage2, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2)
    assert sec_cont.shape == (B, K_MAX, 4)
    assert sec_valid.shape == (B, K_MAX)

@@ -1261,9 +1261,7 @@ def test_build_models_routed_pair_is_drop_in_for_sample_flow():
    n_sec_pred = stage2.predict_n_sec(cond_cont, cond_cat, stage1_norm).argmax(dim=-1)
    assert n_sec_pred.shape == (B,)

-    sec_cont, sec_type_emb, sec_valid = sample_secondaries(
-        stage2, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2
-    )
+    sec_cont, _, sec_valid = sample_secondaries(stage2, cond_cont, cond_cat, stage1_norm, n_sec_pred, steps=2)
    assert sec_cont.shape == (B, K_MAX, 4)
    assert sec_valid.shape == (B, K_MAX)

@@ -262,7 +262,7 @@ def test_sample_secondaries_ar_first_slot_has_no_history():
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
    n_sec_pred = torch.tensor([0, 1, 1])
-    sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
+    sec_cont, _, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, n_sec_pred, steps=2)
    assert sec_cont.shape == (B, 1, CONT_SLOT_DIM)
    assert sec_valid.tolist() == [[False], [True], [True]]

@@ -281,7 +281,7 @@ def test_sample_secondaries_ar_stop_token_forced_stop_gives_zero_secondaries(n_s
    _force_stop_head_logit(decoder, 50.0)
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
-    sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
+    _, _, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
    assert sec_valid.shape == (B, k_max)
    assert not sec_valid.any()

@@ -296,7 +296,7 @@ def test_sample_secondaries_ar_stop_token_forced_never_stop_runs_to_k_max(n_sec_
    _force_stop_head_logit(decoder, -50.0)
    cond_cont, cond_cat = _cond(B)
    stage1_out = torch.randn(B, X_DIM)
-    sec_cont, sec_type, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
+    _, _, sec_valid = sample_secondaries_ar(decoder, cond_cont, cond_cat, stage1_out, None, steps=2)
    assert sec_valid.all()


@@ -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_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) ──────────


@@ -12,6 +12,7 @@ import pytest
 import torch
 import torch.nn.functional as F

+from giant.checkpoint_io import load_for_inference
 from giant.config import ParticleTypeConfig
 from giant.constants import (
    COND_DIM,
@@ -21,7 +22,6 @@ from giant.constants import (
    SEC_SLOT_DIM,
    X_DIM,
 )
-from giant.checkpoint_io import load_for_inference
 from giant.data.dataset import StepBatch
 from giant.data.transforms import Normalizer
 from giant.model.network import Stage2Autoregressive, build_critics, build_models
@@ -36,7 +36,6 @@ from giant.training import (
    train,
 )
 from giant.training.metrics import _wandb_run_config
-from giant.training.trainers import _type_class_weight_vector
 from giant.training.stage2_inputs import (
    _ar_has_prev,
    _assemble_stage2_ar_inputs,
@@ -51,6 +50,7 @@ from giant.training.stage2_inputs import (
    _stop_target_and_mask,
    _type_repr,
 )
+from giant.training.trainers import _type_class_weight_vector

 PDG_VOCAB = 6
 MAT_VOCAB = 3
@@ -543,19 +543,18 @@ def test_train_raises_when_no_active_stage():
    model_config = _model_config(cfg)
    models = build_models(model_config)
    critics = build_critics(model_config)
-    with tempfile.TemporaryDirectory() as tmp:
-        with pytest.raises(ValueError, match="no active stage"):
-            train(
-                cfg=cfg,
-                models=models,
-                critics=critics,
-                train_loader=_fake_batches(1, 8),
-                val_loader=_fake_batches(1, 8),
-                device=torch.device("cpu"),
-                out_dir=Path(tmp) / "run",
-                model_config=model_config,
-                total_train_batches=1,
-            )
+    with tempfile.TemporaryDirectory() as tmp, pytest.raises(ValueError, match="no active stage"):
+        train(
+            cfg=cfg,
+            models=models,
+            critics=critics,
+            train_loader=_fake_batches(1, 8),
+            val_loader=_fake_batches(1, 8),
+            device=torch.device("cpu"),
+            out_dir=Path(tmp) / "run",
+            model_config=model_config,
+            total_train_batches=1,
+        )


 def test_metrics_csv_columns_are_stage_prefixed():
@@ -11,8 +11,8 @@ import pytest

 pytest.importorskip("plotstyle")

-from giant.training import plots as plots_mod  # noqa: E402
-from giant.training.plots import MetricsTable, derive_metrics_dir, render_metrics  # noqa: E402
+from giant.training import plots as plots_mod
+from giant.training.plots import MetricsTable, derive_metrics_dir, render_metrics

 # --- fixtures ----------------------------------------------------------

@@ -2,9 +2,13 @@ import warnings

 import numpy as np
 import pytest
+
 from giant.cond_layout import AXIS_TYPES, CondLayout
 from giant.constants import COND_DIM, COND_DIM_BASE, K_MAX
 from giant.data.transforms import (
+    Normalizer,
+    _vectorized_map_lookup,
+    _WelfordAccumulator,
    build_cond_features,
    build_features,
    encode_secondaries,
@@ -14,12 +18,9 @@ from giant.data.transforms import (
    inv_log_transform,
    local_frame_rotation,
    log_transform,
-    Normalizer,
    reconstruct_post_pos,
    sorted_membership,
    travel_direction,
-    _vectorized_map_lookup,
-    _WelfordAccumulator,
 )
Author	SHA1	Message	Date
lars	ecd6347fca	chore: raise pyarrow ceiling to <26, bump to 25.0.1 CI / Sync project version with tag (hand-pushed tags only) (pull_request) Has been skipped Details CI / Publish package to Gitea package registry (pull_request) Has been skipped Details CI / Type check (ty) (pull_request) Successful in 2m28s Details CI / Lint (ruff check) (pull_request) Successful in 2m28s Details CI / Format (ruff format) (pull_request) Successful in 2m28s Details CI / Tests (pull_request) Successful in 6m2s Details CI / Release (bump, changelog, badges, tag) on merge to master (pull_request) Has been skipped Details 25.0.0's changelog has no breaking changes for our usage (streaming parquet reads via polars/pyarrow) — the only Python-relevant deprecation is the `feather` module, which giant doesn't use. Tests, ruff, and ty all pass against 25.0.1. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TMdZFqXXig7i3XkirSUxef	2026-09-04 14:12:55 +02:00
lars	c984d0a19d	chore: bump uv.lock and fix ruff 0.16 default-rule lint findings CI / Sync project version with tag (hand-pushed tags only) (pull_request) Has been skipped Details CI / Publish package to Gitea package registry (pull_request) Has been skipped Details CI / Format (ruff format) (pull_request) Successful in 53s Details CI / Type check (ty) (pull_request) Successful in 57s Details CI / Lint (ruff check) (pull_request) Successful in 41s Details CI / Tests (pull_request) Successful in 8m20s Details CI / Release (bump, changelog, badges, tag) on merge to master (pull_request) Has been skipped Details uv.lock was stale (ty 0.0.50 -> 0.0.78, ruff 0.15 -> 0.16, polars, numpy, typer, wandb, pytest, and others), all within existing pyproject.toml bounds. ruff 0.16 widened its default rule selection, taking this repo from 0 to 274 lint errors under the same config; --fix handled most of it (import sorting, Optional[X] -> X \| None, ...), and the remainder (unused unpacked variables, dict()-as-literal, subprocess.run without explicit check=, a couple of intentional broad excepts/naive datetimes) were fixed or annotated by hand. Also fixes a real type-narrowing gap ty 0.0.78 caught in test_config_consumed_keys.py's `or`-combined isinstance check. torch stays pinned to 2.3.x (deliberate, see CLAUDE.md); pyarrow's <25 ceiling is left as a separate decision. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01TMdZFqXXig7i3XkirSUxef	2026-09-04 14:09:29 +02:00
gitea-actions	600e04f46a	chore: release v0.3.20 CI / Tests (push) Successful in 2m42s Details CI / Lint (ruff check) (push) Successful in 1m54s Details CI / Format (ruff format) (push) Successful in 2m12s Details CI / Type check (ty) (push) Successful in 2m16s Details CI / Sync project version with tag (hand-pushed tags only) (push) Has been skipped Details CI / Publish package to Gitea package registry (push) Has been skipped Details CI / Release (bump, changelog, badges, tag) on merge to master (push) Successful in 12s Details	2026-09-04 06:47:29 +00:00
lars	085b69081b	Merge pull request 'fix: dereference annotated tag to its commit before the push comparison' (#96 ) from ci/fix-annotated-tag-push-check into master CI / Sync project version with tag (hand-pushed tags only) (push) Has been skipped Details CI / Publish package to Gitea package registry (push) Has been skipped Details CI / Format (ruff format) (push) Successful in 1m25s Details CI / Lint (ruff check) (push) Successful in 1m25s Details CI / Type check (ty) (push) Successful in 1m27s Details CI / Tests (push) Successful in 2m41s Details CI / Release (bump, changelog, badges, tag) on merge to master (push) Successful in 1m56s Details Reviewed-on: #96	2026-09-04 08:41:55 +02:00
ci	b04e7be146	fix: dereference annotated tag to its commit before the push comparison CI / Sync project version with tag (hand-pushed tags only) (pull_request) Has been skipped Details CI / Publish package to Gitea package registry (pull_request) Has been skipped Details CI / Type check (ty) (pull_request) Successful in 35s Details CI / Format (ruff format) (pull_request) Successful in 36s Details CI / Lint (ruff check) (pull_request) Successful in 39s Details CI / Tests (pull_request) Successful in 3m26s Details CI / Release (bump, changelog, badges, tag) on merge to master (pull_request) Has been skipped Details release-commit.sh creates release tags with `git tag -a`, so `git rev-parse "$TAG"` returns the annotated tag object's own SHA, not the commit it points at. Comparing that directly against `git rev-parse HEAD` was therefore always false, so the atomic `git push --atomic origin HEAD:master refs/tags/$TAG` never fired — the release job has been pushing the branch alone and silently dropping every release tag (confirmed: v0.3.19 landed on master but was never pushed as a tag). `^{commit}` dereferences the tag to its target commit so the comparison actually matches. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NnotyatakKNS4NLjDbfYw1	2026-09-04 08:36:17 +02:00
gitea-actions	ff17a5c212	chore: release v0.3.19 CI / Format (ruff format) (push) Has been skipped Details CI / Type check (ty) (push) Has been skipped Details CI / Sync project version with tag (hand-pushed tags only) (push) Successful in 1m7s Details CI / Lint (ruff check) (push) Has been skipped Details CI / Tests (push) Has been skipped Details CI / Release (bump, changelog, badges, tag) on merge to master (push) Has been skipped Details CI / Publish package to Gitea package registry (push) Successful in 1m13s Details	2026-09-04 06:25:44 +00:00
lars	fffde48ed8	Merge pull request 'ci: restructure release pipeline into a single atomic release commit' (#95 ) from ci/restructure-release-pipeline into master CI / Sync project version with tag (hand-pushed tags only) (push) Has been skipped Details CI / Publish package to Gitea package registry (push) Has been skipped Details CI / Lint (ruff check) (push) Successful in 22s Details CI / Format (ruff format) (push) Successful in 1m11s Details CI / Type check (ty) (push) Successful in 1m11s Details CI / Tests (push) Successful in 2m32s Details CI / Release (bump, changelog, badges, tag) on merge to master (push) Successful in 59s Details Reviewed-on: #95	2026-09-04 08:20:54 +02:00
ci	8971121167	ci: restructure release pipeline into a single atomic release commit CI / Sync project version with tag (hand-pushed tags only) (pull_request) Has been skipped Details CI / Publish package to Gitea package registry (pull_request) Has been skipped Details CI / Lint (ruff check) (pull_request) Successful in 19s Details CI / Type check (ty) (pull_request) Successful in 16s Details CI / Format (ruff format) (pull_request) Successful in 24s Details CI / Tests (pull_request) Successful in 2m27s Details CI / Release (bump, changelog, badges, tag) on merge to master (pull_request) Has been skipped Details Every merge to master previously produced three separate commits (bump version -> update changelog, tagged here -> update README badges), so the published tag never carried the current release's own README badges, and the badge commit leaked into the next release's changelog since no cliff.toml parser skipped it. - Extract the bump/changelog/badge assembly into .gitea/scripts/release-commit.sh, used by both the merge-to-master path and the hand-pushed-tag sync path, so every tag now points at one complete "chore: release vX.Y.Z" commit. Push the commit and its tag atomically. - sync-version-on-tag now refuses to touch a tag whose commit isn't reachable from master (rather than silently rewriting an unreviewed tree), and builds a proper release commit via the same script when it does need to correct a hand-pushed tag's version. - publish-package now depends only on sync-version-on-tag: since a tag can only pass that guard if its commit is already on master, and master is always fully checked, re-running the lint/type/test matrix on tag pushes was redundant. - Factor the repeated checkout/setup-uv/env/sync steps into a local composite action (.gitea/actions/setup), fix `test`'s `needs` to include ruff-format, and bump actions/upload-artifact to v4. - cliff.toml: skip "chore: release ..." commits from the changelog. Verified by dry-running release-commit.sh against a scratch worktree for all three code paths (patch bump, --no-bump, explicit VERSION sync), confirming idempotency and that the resulting commit carries pyproject.toml, .bumpversion.toml, uv.lock, CHANGELOG.md and README.md together. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01NnotyatakKNS4NLjDbfYw1	2026-09-03 18:29:38 +02:00
gitea-actions	7cd22a77ce	chore: update README badges (version 0.3.18, 1138 tests) CI / Sync project version with tag (push) Has been skipped Details CI / Lint (ruff check) (push) Successful in 4m18s Details CI / Format (ruff format) (push) Successful in 4m14s Details CI / Type check (ty) (push) Successful in 4m18s Details CI / Tests (push) Successful in 5m44s Details CI / Publish package to Gitea package registry (push) Has been skipped Details CI / Bump version, tag, and update changelog on merge to master (push) Successful in 18s Details CI / Update README badges (version, test count) (push) Successful in 1m35s Details	2026-09-03 16:10:56 +00:00
gitea-actions	c613588a70	chore: update changelog for v0.3.18 CI / Tests (push) Successful in 9m29s Details CI / Lint (ruff check) (push) Successful in 2m6s Details CI / Format (ruff format) (push) Successful in 1m54s Details CI / Type check (ty) (push) Successful in 2m25s Details CI / Sync project version with tag (push) Has been skipped Details CI / Publish package to Gitea package registry (push) Has been skipped Details CI / Bump version, tag, and update changelog on merge to master (push) Successful in 10s Details CI / Update README badges (version, test count) (push) Successful in 1m1s Details	2026-09-03 16:08:53 +00:00
gitea-actions	56642ebd2c	chore: bump version 0.3.17 -> 0.3.18	2026-09-03 16:08:45 +00:00
lars	9a03f4552a	Merge pull request 'perf: compact Stage-2 AR inference loop to active rows only' (#94 ) from perf/stage2-ar-inference-compaction into master CI / Sync project version with tag (push) Has been skipped Details CI / Type check (ty) (push) Successful in 34s Details CI / Lint (ruff check) (push) Successful in 1m23s Details CI / Format (ruff format) (push) Successful in 1m23s Details CI / Tests (push) Successful in 3m2s Details CI / Publish package to Gitea package registry (push) Has been skipped Details CI / Bump version, tag, and update changelog on merge to master (push) Successful in 1m48s Details CI / Update README badges (version, test count) (push) Successful in 54s Details Reviewed-on: #94	2026-09-03 18:03:32 +02:00
lars	5c576fa8f3	perf: compact Stage-2 AR inference loop to active rows only CI / Sync project version with tag (pull_request) Has been skipped Details CI / Type check (ty) (pull_request) Successful in 1m2s Details CI / Lint (ruff check) (pull_request) Successful in 1m6s Details CI / Format (ruff format) (pull_request) Successful in 1m6s Details CI / Tests (pull_request) Successful in 2m46s Details CI / Bump version, tag, and update changelog on merge to master (pull_request) Has been skipped Details CI / Publish package to Gitea package registry (pull_request) Has been skipped Details CI / Update README badges (version, test count) (pull_request) Has been skipped Details sample_secondaries_ar ran all k_max=15 slots for every row regardless of each row's own predicted secondary count, even though the baseline checkpoint's rollout measured only 0.382 secondaries/step — so ~97% of stage-2 model calls generated tokens sec_valid then masked away. Compact the loop to the still-active row set at each slot: drop a row the moment its n_sec_pred is exhausted (or, under n_sec.mode="stop_token", the moment its own stop logit fires), so slot k's model calls cost O(active rows) instead of O(B). Exact — rows are independent given their own history — verified by comparing the compacted path against a new full_length=True escape hatch that reproduces the original uncompacted behavior bit-for-bit under deterministic noise. full_length=True is 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 regardless of a row's own count, so training behavior is unchanged. AttentionHistory's KV cache and MarkovHistory's O(1) state are kept aligned to the shrinking active set via a new HistoryEncoder.select_cache / Stage2Autoregressive.select_history_cache. Also fixes a latent bug the refactor surfaced: derived_n_sec (stop-token mode) could be overwritten by a later spurious re-fire of the stop logit on a row that had already stopped; now tracked via an explicit `finished` mask so only the first stop slot is recorded, matching the documented contract. No architecture or checkpoint-format change — every existing v0.3.0 Stage2Autoregressive checkpoint (flow/wgan, markov/attention, head/stop_token) picks up the speedup automatically on its next `giant rollout`/`giant predict`, no retraining needed. Measured (CPU, hidden_dim=512/6 blocks, k_max=15, batch 512, mean n_sec≈0.38 matching the baseline checkpoint's own rollout): 17.6-22.9x fewer wall-clock seconds for the AR loop alone (attention/markov history respectively). Directional only — baseline.toml's GPU inference-cost comment is updated accordingly, flagged stale pending a real rollout re-measurement via eval_cost_per_step. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01HPt7bVLZYFJe5cG6V7ahqC	2026-09-03 17:56:59 +02:00
lars	bf3271f09e	docs: rewrite README, keep version/test badges live via Gitea Actions CI / Sync project version with tag (push) Has been skipped Details CI / Lint (ruff check) (push) Successful in 1m6s Details CI / Format (ruff format) (push) Successful in 1m6s Details CI / Type check (ty) (push) Successful in 1m12s Details CI / Tests (push) Successful in 3m15s Details CI / Publish package to Gitea package registry (push) Has been skipped Details CI / Bump version, tag, and update changelog on merge to master (push) Successful in 28s Details CI / Update README badges (version, test count) (push) Successful in 1m21s Details Rewrite README.md from scratch as a scannable landing page (hero, one mermaid pipeline diagram, quick start, deep detail folded into collapsible sections) instead of the old flat prose dump duplicating CLAUDE.md. Swap the static "CI" badge for a live Gitea Actions status badge, and add an update-badges job to ci.yml that recomputes the version and test-count badges on every push to master and pushes an update only when they actually changed. Co-Authored-By: Claude Sonnet 5 <noreply@anthropic.com> Claude-Session: https://claude.ai/code/session_01JL1hFbhLv5uwjTXqkWTLnH	2026-09-02 15:58:02 +02:00