SPRIG v0.1 — research preview release

.gitattributes

@@ -33,3 +33,6 @@ saved_model/**/* filter=lfs diff=lfs merge=lfs -text
 *.zip filter=lfs diff=lfs merge=lfs -text
 *.zst filter=lfs diff=lfs merge=lfs -text
 *tfevents* filter=lfs diff=lfs merge=lfs -text
+samples.jpg filter=lfs diff=lfs merge=lfs -text
+samples_prompt_bank.jpg filter=lfs diff=lfs merge=lfs -text
+texel_atlas.png filter=lfs diff=lfs merge=lfs -text

DESIGN.md

@@ -0,0 +1,105 @@
+# SPRIG v0.1 — Design & Contracts (single source of truth)
+
+Reference docs: full architecture spec `~/workspace/SPRIG-C-architecture.md`; project plan `~/.claude/plans/ok-now-i-ltierally-calm-sifakis.md` (data/eval/infra details).
+v0.1 = finite-lattice Stage A at 64x64. No continuous cuts, no morphogen, no occlusion masks, no DINO channel.
+
+## 1. Repo layout & file ownership
+
+```
+sprig/data/procgen/{vocab,sampler,captions,render,writer}.py   # owner: PROCGEN agent
+sprig/data/{dataset,embed_t5}.py, sprig/data/clevr/prep.py     # owner: DATA agent
+sprig/dp/{lattice,inside}.py                                   # owner: DP agent
+sprig/model/{gmt,atlas,dl,sprig}.py                            # owner: MODEL agent
+sprig/eval/{baseline_pixmix,tree_metrics,color_checks,probe,prompts,report,monitors}.py  # owner: EVAL agent
+train.py, configs/*.yaml, infra/*.sh                           # owner: HARNESS agent
+tests/test_<module>_*.py                                       # each agent writes tests for its own modules
+```
+Python 3.12, torch. No deps beyond: torch, numpy, pillow, transformers, sentencepiece, protobuf, tensorboard, matplotlib, pyyaml, tqdm, einops, pytest. Everything must run CPU-only for tests (no `.cuda()` hardcoded; device from config).
+
+## 2. Global config (configs/main64.yaml values)
+
+- Canvas 64x64 RGB. Grid stride g=8 px → 8x8 cells.
+- S=1024 nonterminal symbols, R=64 rule components, T_v=256 texels, d=384 model width.
+- Caption: precomputed T5-base token embeddings [L,768], L≤64.
+- Leaf-eligible regions: both sides ≤ 16 px. Regions with any side > 16 px MUST expand. 8x8 regions MUST terminate.
+- Batch 256. AdamW: lr 3e-4 for embedding tables (E_N, E_T, V, W, P_T, cut-type tables), 1e-4 for GMT + atlas renderer + Φ; betas (0.9, 0.95); cosine decay to 10% over 250k steps, 2k warmup; grad clip 1.0; EMA 0.9999 (eval only).
+- Rule-logit temperature τ_ann: 2.0 → 1.0 linearly over first 50k steps (divide rule/termination logits by τ_ann).
+- 10% of steps: caption replaced by null embedding (dataloader-level).
+- bf16 autocast for matmuls; ALL logsumexp/log-domain accumulation in fp32. Emission DL log-scales clamped [-7, 2].
+
+## 3. Lattice (sprig/dp/lattice.py) — DP agent
+
+Region = axis-aligned rect with corners on the 8-px grid, i.e. cell-interval pair. All (C(9,2))²=1296 rects are regions. Precompute once (pure function of canvas/grid config, cache to `.pt`):
+- `regions: int32 [N_reg, 4]` (x0,y0,x1,y1 in px), `region_id` lookup dict.
+- Cuts: for each region, every interior grid line on each axis is a valid cut → children are regions. Global **cut-type vocabulary** for parameter tying: (axis ∈ {H,V}) × (relative offset bucket ∈ 7 buckets, nearest of {1/8..7/8}) = 14 cut types. Per region: `cut_list` of (cut_type_id, child_lo_id, child_hi_id). child_lo = lesser-coordinate child (left for V, top for H) — this ordering is what lets the grammar express left/right, above/below.
+- Level order: by cell area ascending (levels 1..64). Per level: flattened index tensors `parent_ids [M]`, `cut_type [M]`, `child_lo [M]`, `child_hi [M]` for vectorized DP.
+- `leaf_mask [N_reg]` (both sides ≤16), `must_terminate [N_reg]` (8x8), `must_expand [N_reg]` (any side >16), `phi_geom [N_reg, 64]` Fourier features of (log-area, log-aspect, center-x, center-y).
+- Leaf shape groups: leaf regions grouped by (h,w) ∈ {8,16}² for batched emission scoring.
+
+## 4. Model (sprig/model/) — MODEL agent
+
+**gmt.py — GrammarModulationTransformer.** Symbol embeddings E_N [S,d] (queries; no symbol-symbol self-attention). 4 blocks: {cross-attn to projected caption tokens (768→d, key_padding_mask from emb_len), FFN d→4d→d, pre-LN}. Output H [B,S,d]. Heads:
+- `U = H @ W_u` → p(k|A,c) logits [B,S,R].
+- Termination: MLP([H_A ; phi_geom(r)]) → logit per (B, region, S) — computed lazily as MLP over H then dot with a geometry projection: implement as `term_logit[b,r,A] = MLP_h(H)[b,A,:] · MLP_g(phi_geom)[r,:] + bias_A` (factorized, cheap).
+- Cut-type distribution p(s|k,c): component embeddings e_k [R,d] cross-attend once to caption → logits [B,R,14]; per region, mask to valid cut types and renormalize (log-softmax over the masked set). NOTE: multiple concrete cuts in a region can share a cut type — split p(s|k,c) mass uniformly across same-type concrete cuts (add -log(count) correction, precomputed per region).
+- Terminal texel prior: p(T|A,c) = Σ_k p(k|A,c) softmax(P_T[k]) with static P_T [R,T_v].
+- Children: static V, W [R,S]: p(B|k)=softmax(V[k]), p(C|k)=softmax(W[k]). (v0.1: no q(B) coupling.)
+- Illumination field Φ(c): mean-pooled caption → MLP → deconv to [B,8,16,16], bilinear-resampled at leaf positions, FiLM (scale,shift per DL-mean channel group) on emission means.
+
+**atlas.py — TexelAtlas.** E_T [T_v,d]. Renderer: per texel, cross-attn(E_T row → caption, 2 blocks width 256) → seed [256,4,4] → 2× (conv + 2x nearest-upsample) → atlas [B,T_v,40,16,16]. 40 ch = 4 DL components × (1 weight + 3 means + 3 log-scales + 3 channel-coupling coeffs). Per-texel additive **bias grid** [T_v,40,16,16] (trainable table, no renderer) — the resurrection-writable parameterization (F3.3/M1.2).
+- Emission scoring: for each leaf shape group (h,w): resample atlas 16x16 → (h,w) via adaptive avg pool; FiLM by Φ at leaf position; score pixels under 4-comp discretized logistic with RGB coupling (means: μ_R; μ_G+α·R; μ_B+β·R+γ·G), quantized to 256 bins → `ell [B, n_leaf_in_group, T_v]` summed over pixels. fp32 result.
+
+**dl.py** — discretized logistic mixture log-prob (PixelCNN++ style, vectorized, fp32-safe).
+
+**sprig.py — SPRIGModel** (implements contracts):
+- `log_marginal(image_u8, emb, emb_len) -> logZ [B]` — full inside DP (tempered during training via `self.eta`; an `eta=0` flag for reported numbers).
+- `loss(batch) -> (loss, metrics_dict)` — see §6.
+- `map_parse(image, emb, emb_len) -> list[ParseNode]` — Viterbi (max-semiring) + backtrace. `ParseNode{rect, axis, cut_px, symbol, texel, children}` (dataclass in sprig/model/sprig.py; texel/leaf fields None for internal nodes).
+- `posterior_usage(image, emb, emb_len) -> dict(symbol_usage [S], texel_usage [T_v], node_entropy, emit_mag, rule_mag, mean_depth, mean_leaves)` — expected counts via the autograd identity: make per-(region,symbol) termination potentials and per-(region,cut) potentials require grad, grad of logZ w.r.t. them = posterior marginals; node entropy = occupancy-weighted entropy of conditional split posteriors (good enough for the PI controller).
+- `sample(emb, emb_len, seed_struct, seed_material, n) -> (images_u8 [n,64,64,3], trees)` — ancestral, breadth-parallel over the frontier; two `torch.Generator`s: structural draws (term/k/s/B/C) from seed_struct, material draws (texel choice + pixel sampling; pixels = DL means for v0.1 crispness, texel choice sampled) from seed_material. Best-of-K: `sample_bestof(emb, K)` reranks K derivations by joint log p(τ, x̂|c).
+
+## 5. Inside DP (sprig/dp/inside.py) — DP agent
+
+`inside(ell_leaf, term_logits, cut_logits, U_logmix, logV, logW, lattice, temper_kappa) -> beta [B,N_reg,S] fp32, logZ [B]`
+- β(r,A) for leaf-eligible r: logaddexp( log p_term + logsumexp_T(log p(T|A,c) + ell(r,T)/κ(r)), log(1−p_term) + expand_term(r,A) ) with must_terminate/must_expand masks (∓inf).
+- expand_term via level-synchronous sweep, area ascending. Per level, with index tensors:
+  `Bhat[m,k] = logbmm(beta[:, child_lo[m], :], logV.T)`; same for Chat with logW; `comb[m,k] = logsumexp over cuts grouped by parent of (log p(s=cut_type|k,c) − log count_correction + Bhat + Chat)`; `expand[parent,A] = logbmm(comb, U_log.T)` where U_log = log-softmax(U/τ_ann).
+  Group-by-parent logsumexp via `torch.segment_reduce` or index_put with scatter-logsumexp helper (write one: max-shift per segment, scatter_add of exp, log). All fp32.
+- `logbmm(x_log [.., K], w_log [K, M])`: max-shifted: `x_log.max(-1)` shift, exp, matmul (bf16 ok), log, add shifts. Provide exact fp32 fallback for tests.
+- Tempering: κ(r) = max(1, area_px(r)^η), η a module buffer. PI controller (in train.py, every 2k steps): target node entropy band [0.5, 3.0] nats; if H < 0.5: η += 0.05 + 0.1·(0.5−H); if H > 3.0: η −= 0.05; clamp [0, 1.5]. Anneal η→0 linearly over final 20k steps; all REPORTED bpd at η=0.
+- Viterbi: same sweep, max instead of logsumexp, argmax backtrace tables per level.
+- CORRECTNESS TEST (blocking): tiny lattice (16x16 canvas, 8px grid → 9 regions), tiny config (S=3, R=2, T_v=2, random weights), brute-force enumerate ALL derivation trees (recursive over cuts/symbols/texels), sum exact joint probs, compare to inside logZ within 1e-4. Also: Viterbi tree's joint prob ≤ logZ; posterior marginals from autograd sum to expected node counts.
+
+## 6. Loss & regularizers
+
+```
+L = mean_b( −logZ_tempered / (3·64·64) )                      # nats/subpixel
+  + 1.0 · Σ_T max(0, 1/(4·T_v) − texel_usage_T)               # texel under-use hinge (F3.2)
+  + 0.5 · Σ_A max(0, 1/(4·S)  − symbol_usage_A)               # symbol under-use hinge (M1-style)
+```
+usage vectors = batch-mean posterior expected counts, normalized to sum 1 (from `posterior_usage` computed on the training batch — reuse the same graph's grads: compute via one extra backward-free trick or just take grads of logZ before the optimizer backward; simplest correct: compute usage with `torch.autograd.grad(logZ.sum(), potentials, retain_graph=True)`).
+InfoNCE caption-contrast: DEFERRED; add only if caption-swap margin is flat by 50k steps.
+Dead-texel resurrection (train.py, every 2k steps): texels with usage < 0.1/T_v → overwrite bias grid with (a random training-image 16x16 crop converted to DL-mean params, small noise on other channels), perturb E_T row ±ε. Log resurrection count.
+
+## 7. Training curriculum (procedural main run)
+
+- Steps 0–200k on proc2d with tier reweighting: 0–20k tiers (0:.55,1:.35,2:.10,3:0); 20k–60k (.25,.35,.30,.10); >60k (.10,.30,.40,.20) [target mix].
+- η per PI controller; τ_ann 2→1 over 50k; report-eta-0 bpd from 180k; final 20k η→0 anneal.
+- CLEVR fine-tune (separate config): init from main ckpt, lr×0.3, 30–50k steps, 20% proc2d replay.
+- Smoke config (configs/smoke.yaml): S=128, R=16, T_v=32, d=128, batch 32, 200 steps, val_fast 64 — must run CPU-only on Mac AND on klaus-1 GPU.
+
+## 8. Contracts (C1–C4) — binding for all agents
+
+**C1 batch** (from `sprig/data/dataset.py` collate): `{image: u8 [B,64,64,3], emb: f16 [B,L,768], emb_len: i32 [B], tier: i8 [B], idx: i64 [B]}`. Null-caption substitution (10%, train only) in the dataset/loader using `~/data/sprig/t5/null.f16`.
+**C2** `model.log_marginal(image, emb, emb_len) -> logZ [B]` (η=0 honored when `model.eval()` + `report_mode=True`).
+**C3** `model.map_parse(...)`, `model.posterior_usage(...)` as §4.
+**C4** `model.sample(emb, emb_len, seed_struct, seed_material, n)`, `model.sample_bestof(emb, emb_len, K, seed)`.
+Data formats, eval metrics, prompt bank, run infra: per the plan file §Part 1/3/4 (memmaps, `meta.jsonl` GT-tree schema, 32-prompt bank, `scalars.jsonl`, atomic checkpoints, run_forever.sh) — copy schemas exactly from there.
+
+## 9. Milestone gates (in order, each blocking)
+
+- **G0**: full pytest suite green on Mac CPU (incl. DP brute-force test, holdout scan test, determinism hashes, tree-metric self-test).
+- **G1**: overfit 1 image (script `scripts/overfit1.py`): 2k steps tiny config → bpd < 1.0, MAP parse stable across steps.
+- **G2**: overfit 100 images with captions: Δ_c > 0 (caption info gain), texels alive > 50%.
+- **G3**: klaus-1 smoke (30 min real config): >2 steps/s at batch 256 (else profile before main run), dataloader ≥95% synthetic throughput, ckpt kill/resume works, eval callback produces grids/overlays.
+- **G4**: main run launch. Twice-daily monitor checks per plan alarms.

README.md

@@ -0,0 +1,122 @@
+---
+license: mit
+tags:
+  - text-to-image
+  - scene-grammar
+  - probabilistic-grammar
+  - research-preview
+  - novel-architecture
+library_name: sprig
+pipeline_tag: text-to-image
+---
+
+# SPRIG v0.1 — a text-to-image model where images are *derived*, not denoised
+
+**Research preview.** SPRIG (Stochastic Production-Rule Image Grammar) is a
+from-scratch generative architecture that is **not** a diffusion model, not
+autoregressive, not a GAN, not a VAE. A caption modulates the production
+probabilities of a learned **probabilistic scene grammar**; an image is produced
+by a single top-down **derivation** that recursively splits the canvas into
+typed regions, each painted by a learned "texel" material. Training is **exact
+maximum likelihood** — the marginal over *all* derivation trees, computed by an
+inside dynamic program (log-semiring DP). No noise process, no adversary, no
+ELBO, no token ordering.
+
+This is **v0.1 at 64×64**: a proof-of-concept for the mechanism. It is honest
+about what works and what does not (see the scorecard). ~16M trainable
+parameters on top of a frozen T5-base caption encoder.
+
+<p align="center"><img src="samples.jpg" width="360" alt="SPRIG v0.1 samples"></p>
+
+## What SPRIG does differently
+
+| | Diffusion / Flow | Autoregressive | **SPRIG** |
+|---|---|---|---|
+| Generative act | denoise a fixed grid over many steps | predict tokens in an order | **derive a tree**: recursively split the canvas, commit each node once |
+| Latent | noisy image | token prefix | an *unobserved random tree* summed out |
+| Training | denoising / score matching | next-token likelihood | **exact marginal likelihood** via inside DP |
+| Free bonus | — | — | a real **likelihood** + an interpretable **parse** of any image |
+
+Because analysis and synthesis are the *same* grammar run in two directions, the
+model can also **parse** a real image (infer its most likely derivation) — see
+`parses.png`. This is the strongest, most novel capability and it works well.
+
+## Honest scorecard (v0.1, 50k steps, held-out procedural scenes)
+
+Success criteria were fixed in advance. **Passes 1 of 5** — but the failures are
+localized and understood, not diffuse.
+
+| Gate | Target | Result | |
+|---|---|---|---|
+| Likelihood vs. no-grammar baseline | beat by ≥0.15 bpd | **2.66 vs 6.28 bpd** | ✅ crushes it |
+| Caption information gain Δc | ≥ 0.05 | **0.248** | ✅ 5× |
+| Visible-cut parse F1 | ≥ 0.6 | **0.765** | ✅ parsing works |
+| Object-cell parse recall (tier1/2) | ≥ 0.70 / 0.50 | 0.20 / 0.22 | ❌ scenes too busy |
+| Prompt-swap attribute control | ≥ 0.80 | 0.37 | ❌ partial |
+| — size attribute specifically | — | **1.00** | ✅ size binds perfectly |
+| Spatial-relation accuracy | ≥ 0.70 | 0.00 | ❌ |
+| Compositional holdout (unseen combos) | ≥ 0.60 | 0.01 | ❌ |
+| Grammar health (S_eff / alive texels) | ≥256 / ≥50% | 968 / 43% | ⚠️ texels over-pruned |
+
+**What this means.** The architecture's structural claims are proven: it models
+data far better than a no-grammar baseline, routes caption information, recovers
+scene structure by parsing, and (after a targeted fix) paints real objects. The
+open problem is **caption→object binding**: the model can draw objects and binds
+*size* perfectly, but does not yet reliably paint the *specific* object a prompt
+asks for, and places too many per scene. That is a conditioning/architecture
+issue targeted by v0.2 — not a matter of more training. The full story
+(including a failure diagnosis that traced an earlier version's blank object
+vocabulary to an emission-weighting/tempering interaction) is in the code repo.
+
+## Usage
+
+```bash
+pip install torch safetensors transformers pillow
+# get the `sprig` package + this file from the code repo, then:
+python inference.py --prompt "a red circle on a white background" --out out.png
+```
+
+```python
+from inference import load_sprig, sample
+model = load_sprig("sprig-v0.1.safetensors", "config.json")   # ~16M params, CPU-friendly
+img = sample(model, "a green triangle", seed=0)               # PIL.Image, 64x64
+```
+
+The model outputs native **64×64** images (upscale with nearest-neighbor to
+view). It also returns the derivation tree, so you can inspect *why* each region
+was drawn.
+
+## Files
+
+- `sprig-v0.1.safetensors` — EMA-merged inference weights (60.8 MB, fp32, 15.9M params)
+- `config.json` — architecture config + release metadata
+- `inference.py` — minimal load + sample + T5 caption encoding
+- `metrics.json` — full evaluation numbers
+- `samples.jpg`, `texel_atlas.png`, `parses.png` — qualitative outputs
+- `DESIGN.md` — the concrete v0.1 architecture specification
+
+## Training
+
+64×64, 2M procedural compositional scenes (colored shapes with attributes and
+spatial relations, templated dense captions, held-out attribute combinations),
+frozen T5-base captions precomputed. Exact-likelihood objective + closed-form
+grammar-health regularizers. One RTX PRO 6000 Blackwell GPU, ~50k steps.
+Generator: see the companion dataset repo (seeded, deterministic).
+
+## Limitations & intended use
+
+Research artifact for studying grammar-based generation and exact-likelihood
+text-to-image. **Not** a production image generator: 64×64, synthetic domain,
+object binding incomplete. Samples are blocky by construction (axis-aligned
+region splits). MIT licensed — build on it.
+
+## Citation
+
+```bibtex
+@software{sprig_v0_1_2026,
+  title  = {SPRIG: Text-to-Image by a Stochastic Production-Rule Image Grammar (v0.1)},
+  year   = {2026},
+  note   = {Research preview. Images are derived by a probabilistic scene grammar
+            trained by exact marginal likelihood, not denoised.}
+}
+```

config.json

@@ -0,0 +1,23 @@
+{
+  "architecture": "SPRIG",
+  "version": "0.1",
+  "weights_file": "sprig-v0.1.safetensors",
+  "weights_merged": "ema",
+  "trained_steps": 50000,
+  "canvas_px": 64,
+  "grid_px": 8,
+  "model": {
+    "S": 1024,
+    "R": 64,
+    "T_v": 256,
+    "d": 384,
+    "canvas": 64,
+    "grid": 8,
+    "emb_dim": 768,
+    "L_max": 64,
+    "leaf_chunk": 8,
+    "texel_hinge_weight": 0.5,
+    "emission_obj_weight": 12.0
+  },
+  "note": "Inference-only export: EMA-merged weights, buffers eta/tau are re-initialised at load (tau=1.0, eta=0.0)."
+}

inference.py

@@ -0,0 +1,83 @@
+"""Minimal inference for SPRIG v0.1 — load the released safetensors and sample.
+
+Requires the `sprig` package (https://github.com/  -- or the code repo bundled
+with this release) plus torch, safetensors, transformers (for the T5 caption
+encoder). The model itself is ~16M params and runs on CPU.
+
+    python inference.py --weights sprig-v0.1.safetensors --config config.json \
+        --prompt "a red circle on a white background" --out out.png
+
+Programmatic:
+
+    from inference import load_sprig, sample
+    model = load_sprig("sprig-v0.1.safetensors", "config.json")
+    img = sample(model, "a green triangle", seed=0)   # PIL.Image, 64x64
+"""
+from __future__ import annotations
+
+import argparse
+import json
+from pathlib import Path
+
+import torch
+from PIL import Image
+from safetensors.torch import load_file
+
+_T5 = None
+
+
+def _t5_embed(prompt: str, device: str = "cpu"):
+    """Encode a caption with frozen T5-base -> (emb [1,L,768] f16, len [1] i32)."""
+    global _T5
+    if _T5 is None:
+        from transformers import T5EncoderModel, T5TokenizerFast
+        tok = T5TokenizerFast.from_pretrained("google-t5/t5-base")
+        enc = T5EncoderModel.from_pretrained("google-t5/t5-base").eval().to(device)
+        _T5 = (tok, enc)
+    tok, enc = _T5
+    ids = tok(prompt, return_tensors="pt", truncation=True, max_length=64).to(device)
+    with torch.no_grad():
+        h = enc(**ids).last_hidden_state          # [1, L, 768]
+    n = int(ids["attention_mask"].sum())
+    return h[:, :n].to(torch.float16), torch.tensor([n], dtype=torch.int32, device=device)
+
+
+def load_sprig(weights: str, config: str, device: str = "cpu"):
+    from sprig.model.sprig import SPRIGModel, SPRIGConfig
+    meta = json.loads(Path(config).read_text())
+    fields = set(SPRIGConfig.__dataclass_fields__)
+    cfg = SPRIGConfig(**{k: v for k, v in meta.get("model", {}).items() if k in fields})
+    model = SPRIGModel(cfg)
+    model.load_state_dict(load_file(weights), strict=False)
+    model.tau.fill_(1.0)          # deployment temperature
+    model.eta.fill_(0.0)          # untempered (exact) emissions
+    return model.eval().to(device)
+
+
+def sample(model, prompt: str, seed: int = 0, device: str = "cpu") -> Image.Image:
+    emb, ln = _t5_embed(prompt, device)
+    with torch.no_grad():
+        imgs, _trees = model.sample(emb, ln, seed_struct=seed, seed_material=seed, n=1)
+    return Image.fromarray(imgs[0].cpu().numpy().astype("uint8"))
+
+
+def main() -> None:
+    ap = argparse.ArgumentParser()
+    ap.add_argument("--weights", default="sprig-v0.1.safetensors")
+    ap.add_argument("--config", default="config.json")
+    ap.add_argument("--prompt", required=True)
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument("--out", default="out.png")
+    ap.add_argument("--device", default="cpu")
+    ap.add_argument("--upscale", type=int, default=6)
+    args = ap.parse_args()
+    model = load_sprig(args.weights, args.config, args.device)
+    img = sample(model, args.prompt, args.seed, args.device)
+    if args.upscale > 1:
+        img = img.resize((64 * args.upscale, 64 * args.upscale), Image.NEAREST)
+    img.save(args.out)
+    print("saved", args.out)
+
+
+if __name__ == "__main__":
+    main()

metrics.json

@@ -0,0 +1,78 @@
+{
+  "metrics": {
+    "bpd_val": 2.6107561223037,
+    "bpd_per_tier": {
+      "0": 2.2038232899707055,
+      "1": 2.3735981695328294,
+      "2": 2.6973096096747566,
+      "3": 3.006890315533517
+    },
+    "bpd_tier_ge1": 2.656977891677827,
+    "n": 2000,
+    "delta_c": 0.24781480995112007,
+    "caption_swap_win_frac": 0.25,
+    "b0_bpd": 6.28435796590519,
+    "tree": {
+      "visible_cut_f1": 0.7648812746568768,
+      "leaf_ari": 0.5688396790527814,
+      "recall_per_tier": {
+        "0": 0.19008264462809918,
+        "3": 0.4357142857142857,
+        "1": 0.20238095238095238,
+        "2": 0.22413333333333335
+      },
+      "recall_tier1": 0.20238095238095238,
+      "recall_tier2": 0.22413333333333335
+    },
+    "prompt_swap": {
+      "per_pair": {
+        "a red circle -> a blue circle": 0.0625,
+        "a green square -> a purple square": 0.015625,
+        "a yellow diamond -> a cyan diamond": 0.921875,
+        "an orange ring -> a magenta ring": 0.078125,
+        "a red circle to the left of a blue square -> a red circle to the right of a blue square": 0.0,
+        "a green circle above a yellow square -> a green circle below a yellow square": 0.0,
+        "a small purple square -> a large purple square": 1.0,
+        "a red circle -> a red square": 0.15625
+      },
+      "attribute_move": 0.3723958333333333,
+      "relation_accuracy": 0.0
+    },
+    "holdout_probe_acc": 0.01171875,
+    "heldout_per_combo": {
+      "a blue triangle": 0.03125,
+      "a red ring": 0.015625,
+      "a green star": 0.0,
+      "a yellow cross": 0.0
+    },
+    "health": {
+      "S": 1024,
+      "T_v": 256,
+      "s_eff": 968.4937738012005,
+      "alive_texel_frac": 0.4296875,
+      "node_entropy": 0.13297630846500397
+    }
+  },
+  "gates": {
+    "1_likelihood": {
+      "status": "PASS",
+      "detail": "bpd(tier>=1)=2.657 vs B0=6.284 (need margin >= 0.15); delta_c=0.248 (need >= 0.05)"
+    },
+    "2_parses": {
+      "status": "FAIL",
+      "detail": "recall t1=0.20 (>= 0.7), t2=0.22 (>= 0.5), cut F1=0.76 (>= 0.6)"
+    },
+    "3_prompt_control": {
+      "status": "FAIL",
+      "detail": "attribute-move=0.37 (>= 0.8), relation=0.00 (>= 0.7)"
+    },
+    "4_compositional": {
+      "status": "FAIL",
+      "detail": "held-out combo probe acc=0.01 (>= 0.6)"
+    },
+    "5_health": {
+      "status": "FAIL",
+      "detail": "S_eff=968 (>= 256), alive texels=43% (>= 50%)"
+    }
+  }
+}

requirements.txt

@@ -0,0 +1,6 @@
+torch
+safetensors
+transformers
+pillow
+numpy
+einops

sprig-v0.1.safetensors

@@ -0,0 +1,3 @@
+version https://git-lfs.github.com/spec/v1
+oid sha256:b7ec1c3e062a9769be6dc0d28a2e078b2cf30bc2a4f351bc350a0d191a30a5f0
+size 63701808

sprig/data/clevr/prep.py

@@ -0,0 +1,229 @@
+"""CLEVR v1.0 -> SPRIG memmap preprocessing.
+
+CLI:
+
+    python -m sprig.data.clevr.prep --clevr-root ~/data/CLEVR_v1.0 \
+        --split train --out ~/data/sprig/clevr/train [--limit N] [--seed 0]
+
+Reads `<clevr-root>/scenes/CLEVR_<split>_scenes.json` and
+`<clevr-root>/images/<split>/<image_filename>` (480x320 renders). Each image
+is center-cropped to x in [80, 400) (full height) -> 320x320, LANCZOS-resized
+to 64x64, and written to the standard SPRIG memmap layout (see
+`sprig/data/dataset.py`). Objects whose `pixel_coords` x-center falls outside
+the crop are dropped from the caption pool; the dropped fraction is logged.
+
+Three caption variants are synthesized per image from the scene graph
+(deterministic per (seed, idx)):
+
+  0. pairwise relation read off pixel_coords of two sampled visible objects
+     ("a large red rubber cube to the left of a small blue metal sphere";
+     left/right from pixel x, in front of/behind from pixel y when the
+     x-separation is small);
+  1. partial enumeration capped at 4 objects, order shuffled
+     ("a scene with a ..., a ..., a ... and a ...");
+  2. count + existence ("a scene with five objects, including a gray cube").
+
+Output files: images.u8 [N,64,64,3], meta.jsonl (records with
+"captions": [c0, c1, c2], visible objects, drop counts), meta_offsets.i64
+[N+1] byte offsets, tier_idx/tier0.i64 (CLEVR has no tiers; everything is
+tier 0). Embeddings are computed afterwards by
+`python -m sprig.data.embed_t5 --data-dir <out>`, which detects the
+multi-caption meta and writes emb0/1/2.f16 with their own offsets; the
+dataset then picks one variant per visit (uniform in training).
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+from PIL import Image
+
+CROP_X0 = 80
+CROP_X1 = 400
+IMG_SIZE = 64
+N_CAPTIONS = 3
+# Min horizontal pixel separation (in original 480-px coords) for a
+# left/right relation; below this we use depth-axis (pixel y) instead.
+X_REL_THRESH = 24.0
+
+NUM_WORDS = [
+    "zero", "one", "two", "three", "four", "five",
+    "six", "seven", "eight", "nine", "ten",
+]
+
+
+def num_word(n: int) -> str:
+    return NUM_WORDS[n] if 0 <= n < len(NUM_WORDS) else str(n)
+
+
+def obj_phrase(o: Dict) -> str:
+    return "a %s %s %s %s" % (o["size"], o["color"], o["material"], o["shape"])
+
+
+def crop_resize(img: Image.Image) -> np.ndarray:
+    """480x320 CLEVR render -> center crop x in [80,400) -> 64x64 u8 RGB."""
+    img = img.convert("RGB").crop((CROP_X0, 0, CROP_X1, img.height))
+    img = img.resize((IMG_SIZE, IMG_SIZE), Image.LANCZOS)
+    return np.asarray(img, dtype=np.uint8)
+
+
+def visible_objects(scene: Dict) -> Tuple[List[Dict], int]:
+    """Objects whose pixel-x center lies inside the crop; also #dropped."""
+    kept: List[Dict] = []
+    dropped = 0
+    for o in scene["objects"]:
+        x = float(o["pixel_coords"][0])
+        if CROP_X0 <= x < CROP_X1:
+            kept.append(o)
+        else:
+            dropped += 1
+    return kept, dropped
+
+
+def _relation_caption(objects: List[Dict], rng: np.random.Generator) -> str:
+    if not objects:
+        return "an empty scene"
+    if len(objects) == 1:
+        return obj_phrase(objects[0])
+    i, j = rng.choice(len(objects), size=2, replace=False)
+    a, b = objects[int(i)], objects[int(j)]
+    ax, ay = float(a["pixel_coords"][0]), float(a["pixel_coords"][1])
+    bx, by = float(b["pixel_coords"][0]), float(b["pixel_coords"][1])
+    if abs(ax - bx) >= X_REL_THRESH:
+        rel = "to the left of" if ax < bx else "to the right of"
+    else:
+        # Larger pixel y = lower in frame = closer to the camera.
+        rel = "in front of" if ay > by else "behind"
+    return "%s %s %s" % (obj_phrase(a), rel, obj_phrase(b))
+
+
+def _enumeration_caption(objects: List[Dict], rng: np.random.Generator) -> str:
+    if not objects:
+        return "an empty scene"
+    order = rng.permutation(len(objects))[:4]
+    phrases = [obj_phrase(objects[int(k)]) for k in order]
+    if len(phrases) == 1:
+        listed = phrases[0]
+    else:
+        listed = ", ".join(phrases[:-1]) + " and " + phrases[-1]
+    return "a scene with " + listed
+
+
+def _count_caption(objects: List[Dict], rng: np.random.Generator) -> str:
+    n = len(objects)
+    if n == 0:
+        return "an empty scene"
+    o = objects[int(rng.integers(n))]
+    plural = "object" if n == 1 else "objects"
+    return "a scene with %s %s, including %s" % (num_word(n), plural, obj_phrase(o))
+
+
+def synth_captions(objects: List[Dict], rng: np.random.Generator) -> List[str]:
+    """The three caption variants (relation, enumeration, count/existence)."""
+    return [
+        _relation_caption(objects, rng),
+        _enumeration_caption(objects, rng),
+        _count_caption(objects, rng),
+    ]
+
+
+def prep_split(
+    clevr_root: str,
+    split: str,
+    out_dir: str,
+    limit: Optional[int] = None,
+    seed: int = 0,
+) -> Dict[str, float]:
+    scenes_path = os.path.join(clevr_root, "scenes", "CLEVR_%s_scenes.json" % split)
+    with open(scenes_path, "r") as f:
+        scenes = json.load(f)["scenes"]
+    if limit is not None:
+        scenes = scenes[: int(limit)]
+    n = len(scenes)
+    if n == 0:
+        raise ValueError("no scenes found in %s" % scenes_path)
+
+    os.makedirs(out_dir, exist_ok=True)
+    os.makedirs(os.path.join(out_dir, "tier_idx"), exist_ok=True)
+
+    images = np.memmap(
+        os.path.join(out_dir, "images.u8"),
+        dtype=np.uint8,
+        mode="w+",
+        shape=(n, IMG_SIZE, IMG_SIZE, 3),
+    )
+    meta_offsets = np.zeros(n + 1, dtype=np.int64)
+    total_objects = 0
+    total_dropped = 0
+
+    with open(os.path.join(out_dir, "meta.jsonl"), "wb") as meta_f:
+        for idx, scene in enumerate(scenes):
+            img_path = os.path.join(clevr_root, "images", split, scene["image_filename"])
+            with Image.open(img_path) as img:
+                images[idx] = crop_resize(img)
+
+            kept, dropped = visible_objects(scene)
+            total_objects += len(scene["objects"])
+            total_dropped += dropped
+            rng = np.random.Generator(np.random.PCG64(np.random.SeedSequence([seed, idx])))
+            captions = synth_captions(kept, rng)
+
+            record = {
+                "idx": idx,
+                "image_filename": scene["image_filename"],
+                "tier": 0,
+                "captions": captions,
+                "n_objects": len(kept),
+                "n_dropped": dropped,
+                "objects": [
+                    {
+                        "shape": o["shape"],
+                        "color": o["color"],
+                        "size": o["size"],
+                        "material": o["material"],
+                        "pixel_coords": o["pixel_coords"],
+                    }
+                    for o in kept
+                ],
+            }
+            line = (json.dumps(record) + "\n").encode("utf-8")
+            meta_f.write(line)
+            meta_offsets[idx + 1] = meta_offsets[idx] + len(line)
+
+    images.flush()
+    del images
+    meta_offsets.tofile(os.path.join(out_dir, "meta_offsets.i64"))
+    np.arange(n, dtype=np.int64).tofile(os.path.join(out_dir, "tier_idx", "tier0.i64"))
+
+    drop_frac = total_dropped / max(1, total_objects)
+    stats = {
+        "n_images": float(n),
+        "total_objects": float(total_objects),
+        "total_dropped": float(total_dropped),
+        "drop_frac": drop_frac,
+    }
+    print(
+        "prep %s: %d images; dropped %d/%d objects outside crop (%.2f%%)"
+        % (split, n, total_dropped, total_objects, 100.0 * drop_frac),
+        flush=True,
+    )
+    return stats
+
+
+def main(argv: Optional[List[str]] = None) -> None:
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("--clevr-root", required=True, help="CLEVR_v1.0 directory")
+    ap.add_argument("--split", choices=["train", "val"], default="train")
+    ap.add_argument("--out", required=True, help="output dataset directory")
+    ap.add_argument("--limit", type=int, default=None)
+    ap.add_argument("--seed", type=int, default=0)
+    args = ap.parse_args(argv)
+    prep_split(args.clevr_root, args.split, args.out, args.limit, args.seed)
+
+
+if __name__ == "__main__":
+    main()

sprig/data/dataset.py

@@ -0,0 +1,336 @@
+"""SPRIG dataset over precomputed memmap directories (contract C1).
+
+On-disk layout of a dataset directory (one directory per split), written by
+the procgen writer (`sprig/data/procgen/writer.py`) or by CLEVR prep
+(`sprig/data/clevr/prep.py`) + `sprig/data/embed_t5.py`:
+
+    images.u8            uint8   raw memmap, [N, 64, 64, 3]
+    emb.f16              float16 packed-ragged token embeddings, [total_tokens, 768]
+    emb_offsets.i64      int64   [N+1]; caption i occupies rows offsets[i]:offsets[i+1]
+    meta.jsonl           one JSON object per sample (caption(s), tier, objects, GT tree, ...)
+    meta_offsets.i64     int64   [N+1] byte offsets of line starts into meta.jsonl
+                                 (a trailing entry equal to the file size; an [N]-shaped
+                                 file of line starts is also accepted)
+    tier_idx/tier{t}.i64 int64   sample indices belonging to tier t (t = 0..3)
+
+Multi-caption variant (CLEVR: 3 synthesized captions per image): instead of a
+single `emb.f16`/`emb_offsets.i64` pair, the directory holds one ragged pair
+per caption variant:
+
+    emb0.f16 / emb0_offsets.i64
+    emb1.f16 / emb1_offsets.i64
+    emb2.f16 / emb2_offsets.i64
+
+and each meta line stores `"captions": [c0, c1, c2]`. The dataset picks a
+variant uniformly at random per visit when `train=True`, and deterministically
+as `(idx + epoch) % n_variants` when `train=False` (use `set_epoch` to rotate).
+
+Batch contract C1 (produced by `collate`):
+    {image: u8 [B,64,64,3], emb: f16 [B,Lmax,768] zero-padded,
+     emb_len: i32 [B], tier: i8 [B], idx: i64 [B]}
+
+Null-caption substitution: with probability `p_null` (train only) the caption
+embedding is replaced by the precomputed empty-string embedding loaded from
+`null_emb_path` (packed f16, shape [L0, 768], written by
+`embed_t5.py --null-out`).
+"""
+
+from __future__ import annotations
+
+import json
+import math
+import os
+from typing import Dict, Iterator, List, Optional, Sequence, Tuple
+
+import numpy as np
+import torch
+from torch.utils.data import Dataset, Sampler
+
+EMB_DIM = 768
+IMG_SIZE = 64
+MAX_TIERS = 4
+
+
+def load_null_emb(path: str) -> np.ndarray:
+    """Load the packed-f16 null (empty caption) embedding -> [L0, 768]."""
+    arr = np.fromfile(path, dtype=np.float16)
+    if arr.size == 0 or arr.size % EMB_DIM != 0:
+        raise ValueError("null embedding file %r has invalid size %d" % (path, arr.size))
+    return arr.reshape(-1, EMB_DIM)
+
+
+def load_tier_indices(root: str, n: Optional[int] = None) -> List[np.ndarray]:
+    """Load per-tier index arrays from root/tier_idx/tier{t}.i64.
+
+    If the directory is absent, returns a single tier containing all indices
+    (requires `n`).
+    """
+    tier_dir = os.path.join(root, "tier_idx")
+    if not os.path.isdir(tier_dir):
+        if n is None:
+            raise FileNotFoundError("no tier_idx/ in %r and n not given" % root)
+        return [np.arange(n, dtype=np.int64)]
+    tiers: List[np.ndarray] = []
+    for t in range(MAX_TIERS):
+        p = os.path.join(tier_dir, "tier%d.i64" % t)
+        if os.path.exists(p):
+            tiers.append(np.fromfile(p, dtype=np.int64))
+        elif t == 0:
+            tiers.append(np.zeros(0, dtype=np.int64))
+        else:
+            break
+    return tiers
+
+
+def read_meta(root: str, idx: int) -> Dict:
+    """Random-access read of one meta.jsonl record using meta_offsets.i64."""
+    offsets = np.fromfile(os.path.join(root, "meta_offsets.i64"), dtype=np.int64)
+    with open(os.path.join(root, "meta.jsonl"), "rb") as f:
+        f.seek(int(offsets[idx]))
+        line = f.readline()
+    return json.loads(line)
+
+
+class SprigDataset(Dataset):
+    """Dataset over a memmap directory; see module docstring for the layout."""
+
+    def __init__(
+        self,
+        root: str,
+        p_null: float = 0.0,
+        null_emb_path: Optional[str] = None,
+        train: bool = True,
+        seed: int = 0,
+        emit_obj_mask: bool = False,
+    ) -> None:
+        self.root = root
+        self.train = train
+        self.p_null = float(p_null)
+        self.seed = seed
+        # Object-pixel masks from GT bboxes in meta.jsonl (for the weighted
+        # emission loss / object-crop resurrection). Offsets loaded eagerly;
+        # the meta file handle is opened lazily per worker process.
+        self.emit_obj_mask = bool(emit_obj_mask)
+        self._meta_offsets: Optional[np.ndarray] = None
+        self._meta_file = None
+        if self.emit_obj_mask:
+            self._meta_offsets = np.fromfile(
+                os.path.join(root, "meta_offsets.i64"), dtype=np.int64)
+
+        flat = np.memmap(os.path.join(root, "images.u8"), dtype=np.uint8, mode="r")
+        if flat.size % (IMG_SIZE * IMG_SIZE * 3) != 0:
+            raise ValueError("images.u8 in %r has size not divisible by 64*64*3" % root)
+        self.images = flat.reshape(-1, IMG_SIZE, IMG_SIZE, 3)
+        self.n = self.images.shape[0]
+
+        # Embedding variants: single (emb.f16) or multi (emb0.f16, emb1.f16, ...).
+        pairs: List[Tuple[str, str]] = []
+        if os.path.exists(os.path.join(root, "emb.f16")):
+            pairs.append(("emb.f16", "emb_offsets.i64"))
+        else:
+            v = 0
+            while os.path.exists(os.path.join(root, "emb%d.f16" % v)):
+                pairs.append(("emb%d.f16" % v, "emb%d_offsets.i64" % v))
+                v += 1
+        if not pairs:
+            raise FileNotFoundError("no emb.f16 or emb0.f16 found in %r" % root)
+        self.emb: List[np.ndarray] = []
+        self.emb_offsets: List[np.ndarray] = []
+        for emb_name, off_name in pairs:
+            off = np.fromfile(os.path.join(root, off_name), dtype=np.int64)
+            if off.shape[0] != self.n + 1:
+                raise ValueError(
+                    "%s has %d entries, expected N+1=%d" % (off_name, off.shape[0], self.n + 1)
+                )
+            emb = np.memmap(os.path.join(root, emb_name), dtype=np.float16, mode="r")
+            emb = emb.reshape(-1, EMB_DIM)
+            if emb.shape[0] != int(off[-1]):
+                raise ValueError("%s row count %d != offsets[-1]=%d" % (emb_name, emb.shape[0], int(off[-1])))
+            self.emb.append(emb)
+            self.emb_offsets.append(off)
+        self.n_variants = len(self.emb)
+
+        # Per-sample tier from the tier index arrays (default tier 0).
+        self.tier = np.zeros(self.n, dtype=np.int8)
+        tier_dir = os.path.join(root, "tier_idx")
+        if os.path.isdir(tier_dir):
+            for t in range(MAX_TIERS):
+                p = os.path.join(tier_dir, "tier%d.i64" % t)
+                if os.path.exists(p):
+                    idxs = np.fromfile(p, dtype=np.int64)
+                    self.tier[idxs] = t
+
+        self.null_emb: Optional[np.ndarray] = None
+        if null_emb_path is not None:
+            self.null_emb = load_null_emb(null_emb_path)
+        if self.p_null > 0.0 and self.null_emb is None:
+            raise ValueError("p_null > 0 requires null_emb_path")
+
+        self._epoch = 0
+        self._rng: Optional[np.random.Generator] = None
+
+    def set_epoch(self, epoch: int) -> None:
+        """Rotate the deterministic caption-variant choice (eval mode)."""
+        self._epoch = int(epoch)
+
+    def _get_rng(self) -> np.random.Generator:
+        # Created lazily per worker process so forked workers do not share a
+        # RNG state. Null substitution / variant choice are i.i.d. per visit
+        # and are not part of the checkpointable state (the sampler is).
+        if self._rng is None:
+            info = torch.utils.data.get_worker_info()
+            wid = 0 if info is None else info.id
+            ss = np.random.SeedSequence([self.seed, wid, os.getpid()])
+            self._rng = np.random.Generator(np.random.PCG64(ss))
+        return self._rng
+
+    def __len__(self) -> int:
+        return self.n
+
+    def __getitem__(self, idx: int) -> Dict[str, torch.Tensor]:
+        idx = int(idx)
+        image = torch.from_numpy(np.array(self.images[idx]))  # copy off the memmap
+
+        use_null = False
+        if self.train and self.p_null > 0.0:
+            use_null = bool(self._get_rng().random() < self.p_null)
+
+        if use_null:
+            emb_np = self.null_emb
+        else:
+            if self.n_variants == 1:
+                v = 0
+            elif self.train:
+                v = int(self._get_rng().integers(self.n_variants))
+            else:
+                v = (idx + self._epoch) % self.n_variants
+            off = self.emb_offsets[v]
+            lo, hi = int(off[idx]), int(off[idx + 1])
+            emb_np = self.emb[v][lo:hi]
+        emb = torch.from_numpy(np.array(emb_np, dtype=np.float16))  # copy off the memmap
+
+        out = {
+            "image": image,
+            "emb": emb,
+            "emb_len": torch.tensor(emb.shape[0], dtype=torch.int32),
+            "tier": torch.tensor(int(self.tier[idx]), dtype=torch.int8),
+            "idx": torch.tensor(idx, dtype=torch.int64),
+        }
+        if self.emit_obj_mask:
+            out["objmask"] = self._obj_mask(idx)
+        return out
+
+    def _obj_mask(self, idx: int) -> torch.Tensor:
+        """u8 [IMG_SIZE, IMG_SIZE]: 1 inside any GT object bbox, else 0."""
+        if self._meta_file is None:
+            self._meta_file = open(os.path.join(self.root, "meta.jsonl"), "rb")
+        self._meta_file.seek(int(self._meta_offsets[idx]))
+        m = json.loads(self._meta_file.readline())
+        mask = torch.zeros(IMG_SIZE, IMG_SIZE, dtype=torch.uint8)
+        for o in m.get("objects") or []:
+            bb = o.get("bbox") or o.get("cell")
+            if not bb or len(bb) != 4:
+                continue
+            x0 = max(0, int(math.floor(float(bb[0]))))
+            y0 = max(0, int(math.floor(float(bb[1]))))
+            x1 = min(IMG_SIZE, int(math.ceil(float(bb[2]))))
+            y1 = min(IMG_SIZE, int(math.ceil(float(bb[3]))))
+            if x1 > x0 and y1 > y0:
+                mask[y0:y1, x0:x1] = 1
+        return mask
+
+
+def collate(items: Sequence[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
+    """Collate to contract C1; ragged embeddings zero-padded to the batch max."""
+    b = len(items)
+    lens = [int(it["emb_len"]) for it in items]
+    lmax = max(1, max(lens))
+    emb = torch.zeros(b, lmax, EMB_DIM, dtype=torch.float16)
+    for i, it in enumerate(items):
+        li = lens[i]
+        if li > 0:
+            emb[i, :li] = it["emb"]
+    out = {
+        "image": torch.stack([it["image"] for it in items]),
+        "emb": emb,
+        "emb_len": torch.tensor(lens, dtype=torch.int32),
+        "tier": torch.stack([it["tier"] for it in items]),
+        "idx": torch.stack([it["idx"] for it in items]),
+    }
+    if "objmask" in items[0]:
+        out["objmask"] = torch.stack([it["objmask"] for it in items])
+    return out
+
+
+class TierCurriculumSampler(Sampler):
+    """Infinite sampler drawing tiers according to a stepwise weight schedule.
+
+    schedule: list of (step_start, [w0..w_{T-1}]) pairs; the entry with the
+    largest step_start <= current training step is active. `batch_size`
+    converts sample draws into training steps (step = draws // batch_size), so
+    schedules can be written in optimizer steps as in DESIGN.md section 7.
+
+    Checkpointable via state_dict/load_state_dict (RNG state + draw counter).
+    Iterate it from the main process (the default for torch DataLoader
+    samplers) so the state advances where checkpoints are taken; with
+    num_workers > 0 the prefetch queue makes resumes approximate by up to
+    (num_workers * prefetch_factor) batches, which is acceptable.
+    """
+
+    def __init__(
+        self,
+        tier_indices: Sequence[np.ndarray],
+        schedule: Sequence[Tuple[int, Sequence[float]]],
+        batch_size: int = 1,
+        seed: int = 0,
+    ) -> None:
+        self.tier_indices = [np.asarray(a, dtype=np.int64) for a in tier_indices]
+        if not self.tier_indices:
+            raise ValueError("tier_indices is empty")
+        sched = sorted(((int(s), list(map(float, w))) for s, w in schedule), key=lambda x: x[0])
+        if not sched or sched[0][0] != 0:
+            raise ValueError("schedule must start at step 0")
+        for _, w in sched:
+            if len(w) != len(self.tier_indices):
+                raise ValueError("schedule weight vectors must have one entry per tier")
+        self.schedule = sched
+        self.batch_size = int(batch_size)
+        self.seed = seed
+        self._draws = 0
+        self._rng = np.random.Generator(np.random.PCG64(seed))
+
+    def _weights_at(self, step: int) -> np.ndarray:
+        w = self.schedule[0][1]
+        for s0, wi in self.schedule:
+            if s0 <= step:
+                w = wi
+            else:
+                break
+        w = np.asarray(w, dtype=np.float64)
+        # Zero out empty tiers so we never draw from them.
+        for t, arr in enumerate(self.tier_indices):
+            if arr.size == 0:
+                w[t] = 0.0
+        total = w.sum()
+        if total <= 0.0:
+            raise ValueError("all active tiers at step %d are empty or zero-weight" % step)
+        return w / total
+
+    def __iter__(self) -> Iterator[int]:
+        n_tiers = len(self.tier_indices)
+        while True:
+            step = self._draws // self.batch_size
+            w = self._weights_at(step)
+            t = int(self._rng.choice(n_tiers, p=w))
+            arr = self.tier_indices[t]
+            i = int(arr[int(self._rng.integers(arr.size))])
+            self._draws += 1
+            yield i
+
+    def state_dict(self) -> Dict:
+        return {"draws": self._draws, "rng_state": self._rng.bit_generator.state}
+
+    def load_state_dict(self, state: Dict) -> None:
+        self._draws = int(state["draws"])
+        self._rng = np.random.Generator(np.random.PCG64(self.seed))
+        self._rng.bit_generator.state = state["rng_state"]

sprig/data/embed_t5.py

@@ -0,0 +1,303 @@
+"""Precompute frozen T5-base caption embeddings for SPRIG datasets.
+
+CLI (run once per dataset directory; the encoder never runs during training):
+
+    python -m sprig.data.embed_t5 --data-dir ~/data/sprig/proc2d/train
+    python -m sprig.data.embed_t5 --null-out ~/data/sprig/t5/null.f16
+    python -m sprig.data.embed_t5 --prompts-out ~/data/sprig/t5/promptbank.npz \
+        [--prompts-file prompts.json]
+
+--data-dir reads captions from <dir>/meta.jsonl in order and writes
+packed-ragged fp16 token embeddings (valid tokens only, per the tokenizer
+attention mask) plus int64 offsets [N+1]:
+
+  * single-caption meta ({"caption": ...})      -> emb.f16 / emb_offsets.i64
+  * multi-caption meta ({"captions": [c0..ck]}) -> emb{v}.f16 / emb{v}_offsets.i64
+    for every variant v (or just one with --variant V).
+
+Captions are tokenized with truncation to --max-len (64) tokens,
+length-bucketed within --chunk-size (100k) caption chunks, encoded in batches
+of --batch-size (512), and written back in the original caption order.
+
+--null-out embeds the empty string (a single </s> token embedding) to a
+packed f16 file — the null-caption substitute used by the dataloader (C1).
+
+--prompts-out embeds a JSON list of prompts (from --prompts-file, else from
+`sprig.eval.prompts.PROMPTS`, imported lazily) into an npz with keys
+{emb: f16 [P, Lmax, 768] zero-padded, len: i32 [P], prompts: str array}.
+"""
+
+from __future__ import annotations
+
+import argparse
+import json
+import os
+from typing import Callable, Iterator, List, Optional, Sequence, Tuple
+
+import numpy as np
+
+EMB_DIM = 768
+DEFAULT_MODEL = "google-t5/t5-base"
+
+
+def read_meta_captions(meta_path: str) -> Tuple[List[str], int]:
+    """Read captions from meta.jsonl.
+
+    Returns (flat_captions, n_variants): for single-caption meta the flat list
+    is the N captions and n_variants == 1; for multi-caption meta the list is
+    variant-major, i.e. [all c0, then all c1, ...], length N * n_variants.
+    """
+    singles: List[str] = []
+    multis: List[List[str]] = []
+    with open(meta_path, "r") as f:
+        for line in f:
+            line = line.strip()
+            if not line:
+                continue
+            d = json.loads(line)
+            if "captions" in d:
+                multis.append(list(d["captions"]))
+            else:
+                singles.append(d["caption"])
+    if singles and multis:
+        raise ValueError("meta.jsonl mixes 'caption' and 'captions' records")
+    if singles:
+        return singles, 1
+    if not multis:
+        return [], 1
+    n_var = len(multis[0])
+    for caps in multis:
+        if len(caps) != n_var:
+            raise ValueError("inconsistent number of caption variants in meta.jsonl")
+    flat: List[str] = []
+    for v in range(n_var):
+        flat.extend(caps[v] for caps in multis)
+    return flat, n_var
+
+
+def embed_corpus(
+    captions: Sequence[str],
+    encode_batch: Callable[[List[str]], List[np.ndarray]],
+    token_len: Callable[[str], int],
+    batch_size: int = 512,
+    chunk_size: int = 100000,
+) -> Iterator[np.ndarray]:
+    """Yield one [L_i, 768] f16 array per caption, in the original order.
+
+    Within each chunk of `chunk_size` captions, indices are sorted by token
+    length so batches are near-uniform in length (minimal padding waste), then
+    the results are restored to the original order before yielding.
+    """
+    n = len(captions)
+    for c0 in range(0, n, chunk_size):
+        chunk = list(captions[c0 : c0 + chunk_size])
+        lens = np.asarray([token_len(c) for c in chunk], dtype=np.int64)
+        order = np.argsort(lens, kind="stable")
+        results: List[Optional[np.ndarray]] = [None] * len(chunk)
+        for b0 in range(0, len(order), batch_size):
+            batch_idx = order[b0 : b0 + batch_size]
+            embs = encode_batch([chunk[int(j)] for j in batch_idx])
+            if len(embs) != len(batch_idx):
+                raise RuntimeError("encode_batch returned wrong number of embeddings")
+            for j, e in zip(batch_idx, embs):
+                results[int(j)] = np.asarray(e, dtype=np.float16)
+        for r in results:
+            assert r is not None
+            if r.ndim != 2 or r.shape[1] != EMB_DIM:
+                raise RuntimeError("embedding has shape %r, expected [L, %d]" % (r.shape, EMB_DIM))
+            yield r
+
+
+def write_packed(
+    arrays: Iterator[np.ndarray], n: int, emb_path: str, offsets_path: str
+) -> np.ndarray:
+    """Stream n ragged [L,768] f16 arrays to a packed file + [N+1] offsets."""
+    offsets = np.zeros(n + 1, dtype=np.int64)
+    tmp = emb_path + ".tmp"
+    count = 0
+    with open(tmp, "wb") as f:
+        for i, a in enumerate(arrays):
+            if i >= n:
+                raise RuntimeError("more arrays than expected (n=%d)" % n)
+            a.astype(np.float16).tofile(f)
+            offsets[i + 1] = offsets[i] + a.shape[0]
+            count += 1
+    if count != n:
+        os.remove(tmp)
+        raise RuntimeError("expected %d arrays, got %d" % (n, count))
+    os.replace(tmp, emb_path)
+    offsets.tofile(offsets_path)
+    return offsets
+
+
+def _load_t5_encoder_nommap(model_name: str, dtype):
+    """Load T5EncoderModel with weights read via plain file reads (no mmap).
+
+    transformers/safetensors normally mmap the checkpoint and materialize
+    weight pages lazily at forward time; on memory-constrained hosts (e.g. an
+    8 GB Mac under pressure) those page-ins can SIGBUS the process. Reading
+    the bytes up front and deserializing in memory avoids the mmap entirely.
+    Only handles single-file safetensors checkpoints; callers fall back to
+    from_pretrained otherwise.
+    """
+    import torch
+    from huggingface_hub import hf_hub_download
+    from safetensors.torch import load as st_load
+    from transformers import AutoConfig, T5EncoderModel
+
+    cfg = AutoConfig.from_pretrained(model_name)
+    path = hf_hub_download(model_name, "model.safetensors")
+    with open(path, "rb") as f:
+        data = f.read()
+    sd = st_load(data)
+    del data
+    model = T5EncoderModel(cfg)
+    # The checkpoint carries the full T5; keep encoder weights only
+    # (embed_tokens is tied to `shared` at construction).
+    model.load_state_dict(sd, strict=False)
+    model.tie_weights()
+    return model.to(dtype)
+
+
+def make_t5_encoder(
+    model_name: str = DEFAULT_MODEL, device: str = "cpu", max_len: int = 64
+) -> Tuple[Callable[[str], int], Callable[[List[str]], List[np.ndarray]]]:
+    """Build (token_len, encode_batch) over a frozen T5 encoder.
+
+    bf16 on cuda, fp32 on cpu; outputs are cast to fp16 numpy arrays holding
+    valid (attention-masked) token embeddings only.
+    """
+    import torch
+    from transformers import AutoTokenizer, T5EncoderModel
+
+    tok = AutoTokenizer.from_pretrained(model_name)
+    dtype = torch.bfloat16 if device.startswith("cuda") else torch.float32
+    try:
+        model = _load_t5_encoder_nommap(model_name, dtype)
+    except Exception:
+        model = T5EncoderModel.from_pretrained(model_name, dtype=dtype)
+    model = model.to(device).eval()
+
+    def token_len(text: str) -> int:
+        return len(tok(text, truncation=True, max_length=max_len).input_ids)
+
+    def encode_batch(texts: List[str]) -> List[np.ndarray]:
+        enc = tok(
+            texts,
+            truncation=True,
+            max_length=max_len,
+            padding=True,
+            return_tensors="pt",
+        ).to(device)
+        with torch.no_grad():
+            out = model(**enc).last_hidden_state  # [B, L, 768]
+        mask = enc.attention_mask.bool()
+        return [
+            out[i, mask[i]].to(torch.float16).cpu().numpy() for i in range(len(texts))
+        ]
+
+    return token_len, encode_batch
+
+
+def _load_prompts(prompts_file: Optional[str]) -> List[str]:
+    if prompts_file is not None:
+        with open(prompts_file, "r") as f:
+            prompts = json.load(f)
+        if not isinstance(prompts, list):
+            raise ValueError("--prompts-file must contain a JSON list of strings")
+        return [str(p) for p in prompts]
+    from sprig.eval.prompts import PROMPTS  # lazy: eval module may not exist yet
+
+    return list(PROMPTS)
+
+
+def _embed_data_dir(
+    data_dir: str,
+    variant: Optional[int],
+    token_len: Callable[[str], int],
+    encode_batch: Callable[[List[str]], List[np.ndarray]],
+    batch_size: int,
+    chunk_size: int,
+) -> None:
+    captions, n_var = read_meta_captions(os.path.join(data_dir, "meta.jsonl"))
+    n = len(captions) // n_var
+    variants = range(n_var) if variant is None else [variant]
+    for v in variants:
+        if not (0 <= v < n_var):
+            raise ValueError("--variant %d out of range (meta has %d variants)" % (v, n_var))
+        caps_v = captions[v * n : (v + 1) * n]
+        if n_var == 1:
+            emb_path = os.path.join(data_dir, "emb.f16")
+            off_path = os.path.join(data_dir, "emb_offsets.i64")
+        else:
+            emb_path = os.path.join(data_dir, "emb%d.f16" % v)
+            off_path = os.path.join(data_dir, "emb%d_offsets.i64" % v)
+        arrays = embed_corpus(caps_v, encode_batch, token_len, batch_size, chunk_size)
+        offsets = write_packed(arrays, n, emb_path, off_path)
+        print(
+            "wrote %s: %d captions, %d tokens" % (emb_path, n, int(offsets[-1])),
+            flush=True,
+        )
+
+
+def main(argv: Optional[List[str]] = None) -> None:
+    ap = argparse.ArgumentParser(description=__doc__)
+    ap.add_argument("--data-dir", default=None, help="dataset dir with meta.jsonl")
+    ap.add_argument(
+        "--variant",
+        type=int,
+        default=None,
+        help="embed only this caption variant (default: all variants in meta)",
+    )
+    ap.add_argument("--model", default=DEFAULT_MODEL)
+    ap.add_argument("--device", default=None, help="cpu/cuda (default: auto)")
+    ap.add_argument("--batch-size", type=int, default=512)
+    ap.add_argument("--max-len", type=int, default=64)
+    ap.add_argument("--chunk-size", type=int, default=100000)
+    ap.add_argument("--null-out", default=None, help="write empty-caption embedding here")
+    ap.add_argument("--prompts-out", default=None, help="write promptbank.npz here")
+    ap.add_argument("--prompts-file", default=None, help="JSON list of prompts")
+    args = ap.parse_args(argv)
+
+    if args.data_dir is None and args.null_out is None and args.prompts_out is None:
+        ap.error("nothing to do: give --data-dir, --null-out, and/or --prompts-out")
+
+    device = args.device
+    if device is None:
+        import torch
+
+        device = "cuda" if torch.cuda.is_available() else "cpu"
+
+    token_len, encode_batch = make_t5_encoder(args.model, device, args.max_len)
+
+    if args.data_dir is not None:
+        _embed_data_dir(
+            args.data_dir,
+            args.variant,
+            token_len,
+            encode_batch,
+            args.batch_size,
+            args.chunk_size,
+        )
+
+    if args.null_out is not None:
+        null = encode_batch([""])[0]
+        os.makedirs(os.path.dirname(os.path.abspath(args.null_out)), exist_ok=True)
+        null.astype(np.float16).tofile(args.null_out)
+        print("wrote %s: [%d, %d]" % (args.null_out, null.shape[0], null.shape[1]))
+
+    if args.prompts_out is not None:
+        prompts = _load_prompts(args.prompts_file)
+        embs = encode_batch(prompts) if prompts else []
+        lens = np.asarray([e.shape[0] for e in embs], dtype=np.int32)
+        lmax = int(lens.max()) if len(embs) else 1
+        emb = np.zeros((len(prompts), lmax, EMB_DIM), dtype=np.float16)
+        for i, e in enumerate(embs):
+            emb[i, : e.shape[0]] = e
+        os.makedirs(os.path.dirname(os.path.abspath(args.prompts_out)), exist_ok=True)
+        np.savez(args.prompts_out, emb=emb, len=lens, prompts=np.array(prompts))
+        print("wrote %s: %d prompts, Lmax=%d" % (args.prompts_out, len(prompts), lmax))
+
+
+if __name__ == "__main__":
+    main()

sprig/data/procgen/captions.py

@@ -0,0 +1,219 @@
+"""Templated dense captions for procedural scenes.
+
+10 training templates (T1..T10) + 2 held-out eval templates (E1, E2 — never
+sampled in training mode). 15% of training captions are partial (attribute or
+object dropping) and flagged `partial=True` to teach marginalization.
+
+Template inventory (tiers they apply to):
+  T1  attribute                        "a small striped red circle"          (0)
+  T2  relation, subject-first          "a X to the left of a Y"              (1)
+  T3  attribute + background           "a X on a gray background"            (0)
+  T4  enumerative                      "a scene with a X, a Y, and a Z"      (1,2)
+  T5  relation, INVERTED mention order "a Y to the right of a X"             (1)
+  T6  containment                      "a X inside a blue frame"             (3)
+  T7  count + background               "a scene with three shapes on ..."    (1,2)
+  T8  background-first relation        "on a gray background, a X above a Y" (1)
+  T9  count + enumeration              "three shapes: a X, a Y, and a Z"     (2)
+  T10 containment, frame-first         "a blue frame containing a X"         (3)
+  E1  eval rephrase                    "the image shows ..."                 (all)
+  E2  eval rephrase                    "there is ... in the picture"         (all)
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+
+import numpy as np
+
+from .sampler import Scene
+from .vocab import TEXTURE_ADJ
+
+TRAIN_TEMPLATE_IDS: Tuple[str, ...] = (
+    "T1", "T2", "T3", "T4", "T5", "T6", "T7", "T8", "T9", "T10",
+)
+EVAL_TEMPLATE_IDS: Tuple[str, ...] = ("E1", "E2")
+
+_TIER_TEMPLATES: Dict[int, Tuple[str, ...]] = {
+    0: ("T1", "T3"),
+    1: ("T2", "T5", "T4", "T7", "T8"),
+    2: ("T4", "T7", "T9"),
+    3: ("T6", "T10"),
+}
+
+_NUM_WORDS = {1: "one", 2: "two", 3: "three", 4: "four", 5: "five"}
+# relation type -> (forward phrase for A rel B, inverted phrase for B rel A)
+_REL = {"left": ("to the left of", "to the right of"), "above": ("above", "below")}
+
+PARTIAL_RATE = 0.15
+
+
+@dataclass
+class Caption:
+    text: str
+    template_id: str
+    partial: bool
+
+
+def _article(phrase: str) -> str:
+    return "an" if phrase and phrase[0] in "aeiou" else "a"
+
+
+def _obj_phrase(obj: Dict[str, Any], drops: Sequence[str] = ()) -> str:
+    words: List[str] = []
+    if "size" not in drops:
+        words.append(obj["size"])
+    adj = TEXTURE_ADJ[obj["texture"]]
+    if adj and "texture" not in drops:
+        words.append(adj)
+    if "color" not in drops:
+        words.append(obj["color"])
+    words.append(obj["shape"])
+    core = " ".join(words)
+    return "{} {}".format(_article(core), core)
+
+
+def _list_phrase(phrases: Sequence[str]) -> str:
+    if len(phrases) == 1:
+        return phrases[0]
+    if len(phrases) == 2:
+        return "{} and {}".format(phrases[0], phrases[1])
+    return "{}, and {}".format(", ".join(phrases[:-1]), phrases[-1])
+
+
+def _bg_phrase(scene: Scene) -> str:
+    names = scene.background.split("|")
+    if len(names) == 2:
+        return "a {} and {} background".format(names[0], names[1])
+    return "a {} background".format(names[0])
+
+
+def _relation_parts(
+    scene: Scene, drops: Sequence[str]
+) -> Tuple[str, str, str, str]:
+    rel = scene.relation
+    assert rel is not None, "relation template on a scene without a relation"
+    fwd, inv = _REL[rel["type"]]
+    pa = _obj_phrase(scene.objects[rel["a"]], drops)
+    pb = _obj_phrase(scene.objects[rel["b"]], drops)
+    return pa, pb, fwd, inv
+
+
+def render_template(
+    scene: Scene,
+    template_id: str,
+    drops: Sequence[str] = (),
+    obj_order: Optional[Sequence[int]] = None,
+) -> str:
+    """Realize `template_id` for `scene`.
+
+    `drops`: subset of {"size","texture","color","bg"} removed from the text
+    (partial captions). `obj_order`: object index order/subset for the
+    enumerative templates T4/T9/E1/E2.
+    """
+    objs = scene.objects
+    order = list(obj_order) if obj_order is not None else list(range(len(objs)))
+    phrases = [_obj_phrase(objs[i], drops) for i in order]
+
+    if template_id == "T1":
+        return phrases[0]
+    if template_id == "T2":
+        pa, pb, fwd, _ = _relation_parts(scene, drops)
+        return "{} {} {}".format(pa, fwd, pb)
+    if template_id == "T3":
+        if "bg" in drops:
+            return phrases[0]
+        return "{} on {}".format(phrases[0], _bg_phrase(scene))
+    if template_id == "T4":
+        return "a scene with {}".format(_list_phrase(phrases))
+    if template_id == "T5":
+        pa, pb, _, inv = _relation_parts(scene, drops)
+        return "{} {} {}".format(pb, inv, pa)
+    if template_id == "T6":
+        assert scene.frame is not None
+        fc = scene.frame["color"]
+        return "{} inside {} {} frame".format(phrases[0], _article(fc), fc)
+    if template_id == "T7":
+        n = _NUM_WORDS[len(objs)]
+        noun = "shape" if len(objs) == 1 else "shapes"
+        if "bg" in drops:
+            return "a scene with {} {}".format(n, noun)
+        return "a scene with {} {} on {}".format(n, noun, _bg_phrase(scene))
+    if template_id == "T8":
+        pa, pb, fwd, _ = _relation_parts(scene, drops)
+        if "bg" in drops:
+            return "{} {} {}".format(pa, fwd, pb)
+        return "on {}, {} {} {}".format(_bg_phrase(scene), pa, fwd, pb)
+    if template_id == "T9":
+        n = _NUM_WORDS[len(order)]
+        noun = "shape" if len(order) == 1 else "shapes"
+        return "{} {}: {}".format(n, noun, _list_phrase(phrases))
+    if template_id == "T10":
+        assert scene.frame is not None
+        fc = scene.frame["color"]
+        return "{} {} frame containing {}".format(_article(fc), fc, phrases[0])
+    if template_id in ("E1", "E2"):
+        content = _eval_content(scene, phrases)
+        if template_id == "E1":
+            return "the image shows {}".format(content)
+        verb = "are" if len(objs) > 1 else "is"
+        return "there {} {} in the picture".format(verb, content)
+    raise ValueError("unknown template id: {}".format(template_id))
+
+
+def _eval_content(scene: Scene, phrases: Sequence[str]) -> str:
+    if scene.tier == 1 and scene.relation is not None:
+        pa, pb, fwd, _ = _relation_parts(scene, ())
+        return "{} {} {}".format(pa, fwd, pb)
+    if scene.tier == 3 and scene.frame is not None:
+        fc = scene.frame["color"]
+        return "{} inside {} {} frame".format(phrases[0], _article(fc), fc)
+    return _list_phrase(phrases)
+
+
+def _drop_candidates(scene: Scene, template_id: str) -> List[str]:
+    cands: List[str] = []
+    if template_id != "T7":  # T7 mentions no object attributes
+        cands.extend(["size", "color"])
+        if any(TEXTURE_ADJ[o["texture"]] for o in scene.objects):
+            cands.append("texture")
+    if template_id in ("T3", "T7", "T8"):
+        cands.append("bg")
+    return cands
+
+
+def sample_caption(
+    scene: Scene, rng: np.random.Generator, mode: str = "train"
+) -> Caption:
+    """Sample a caption for `scene`.
+
+    mode="train": tier-appropriate training templates, 15% partial captions.
+    mode="eval": held-out E1/E2 templates only, never partial.
+    """
+    if mode == "eval":
+        tid = EVAL_TEMPLATE_IDS[int(rng.integers(len(EVAL_TEMPLATE_IDS)))]
+        return Caption(render_template(scene, tid), tid, False)
+    assert mode == "train", "mode must be 'train' or 'eval'"
+
+    pool = _TIER_TEMPLATES[scene.tier]
+    tid = pool[int(rng.integers(len(pool)))]
+
+    obj_order: Optional[List[int]] = None
+    if tid in ("T4", "T9"):
+        obj_order = [int(i) for i in rng.permutation(len(scene.objects))]
+
+    partial = bool(rng.random() < PARTIAL_RATE)
+    drops: Tuple[str, ...] = ()
+    if partial:
+        if tid == "T4" and len(scene.objects) >= 2 and rng.random() < 0.5:
+            # object dropping: keep a strict nonempty subset
+            keep = 1 + int(rng.integers(len(scene.objects) - 1))
+            assert obj_order is not None
+            obj_order = obj_order[:keep]
+        else:
+            cands = _drop_candidates(scene, tid)
+            k = 1 if (len(cands) == 1 or rng.random() < 0.5) else 2
+            picked = rng.permutation(len(cands))[:k]
+            drops = tuple(cands[int(i)] for i in picked)
+
+    text = render_template(scene, tid, drops=drops, obj_order=obj_order)
+    return Caption(text, tid, partial)

sprig/data/procgen/render.py

@@ -0,0 +1,111 @@
+"""Deterministic rasterizer for procedural scenes.
+
+Draws at 256x256 (4x supersample) with numpy + PIL, then BOX-downsamples to
+64x64. `render_scene` is a pure function of the Scene, and the Scene is a pure
+function of (global_seed, idx) via SeedSequence([global_seed, idx]) — so the
+same idx always yields bit-identical pixels (regression-tested by hash).
+"""
+from __future__ import annotations
+
+import math
+from typing import Any, Dict, List, Optional, Sequence, Tuple
+
+import numpy as np
+from PIL import Image, ImageDraw
+
+from .sampler import DEFAULT_TIER_MIX, Scene, iter_leaves, sample_scene
+from .vocab import CANVAS, SUPERSAMPLE
+
+_S = SUPERSAMPLE
+_HI = CANVAS * _S
+_BOX = getattr(getattr(Image, "Resampling", Image), "BOX")
+
+_texture_cache: Dict[str, np.ndarray] = {}
+
+
+def _texture_mask(name: str) -> np.ndarray:
+    """Boolean [256,256] pattern mask in canvas-aligned hi-res coordinates."""
+    if name in _texture_cache:
+        return _texture_cache[name]
+    yy, xx = np.mgrid[0:_HI, 0:_HI]
+    if name == "solid":
+        m = np.ones((_HI, _HI), dtype=bool)
+    elif name == "striped":  # diagonal stripes, 2 canvas px on / 2 off
+        m = ((xx + yy) // (2 * _S)) % 2 == 0
+    elif name == "checker":  # 2x2 canvas px checkerboard
+        m = ((xx // (2 * _S)) + (yy // (2 * _S))) % 2 == 0
+    elif name == "dotted":  # dot grid, period 4 canvas px, dot radius ~1.3 px
+        p = 4 * _S
+        dx = (xx % p) - p / 2 + 0.5
+        dy = (yy % p) - p / 2 + 0.5
+        m = dx * dx + dy * dy <= (1.3 * _S) ** 2
+    else:
+        raise ValueError("unknown texture: {}".format(name))
+    _texture_cache[name] = m
+    return m
+
+
+def _shape_mask(shape: str, bbox: Sequence[float]) -> np.ndarray:
+    """Boolean [256,256] mask of `shape` drawn inside hi-res bbox."""
+    x0, y0, x1, y1 = (v * _S for v in bbox)
+    w, h = x1 - x0, y1 - y0
+    cx, cy = (x0 + x1) / 2.0, (y0 + y1) / 2.0
+    im = Image.new("L", (_HI, _HI), 0)
+    d = ImageDraw.Draw(im)
+    if shape == "circle":
+        d.ellipse([x0, y0, x1, y1], fill=255)
+    elif shape in ("square", "rectangle"):
+        d.rectangle([x0, y0, x1, y1], fill=255)
+    elif shape == "triangle":
+        d.polygon([(cx, y0), (x1, y1), (x0, y1)], fill=255)
+    elif shape == "diamond":
+        d.polygon([(cx, y0), (x1, cy), (cx, y1), (x0, cy)], fill=255)
+    elif shape == "star":
+        pts: List[Tuple[float, float]] = []
+        for k in range(10):
+            ang = -math.pi / 2 + k * math.pi / 5
+            r = 1.0 if k % 2 == 0 else 0.42
+            pts.append((cx + math.cos(ang) * r * w / 2, cy + math.sin(ang) * r * h / 2))
+        d.polygon(pts, fill=255)
+    elif shape == "cross":
+        tx, ty = w * 0.34, h * 0.34
+        d.rectangle([cx - tx / 2, y0, cx + tx / 2, y1], fill=255)
+        d.rectangle([x0, cy - ty / 2, x1, cy + ty / 2], fill=255)
+    elif shape == "ring":
+        d.ellipse([x0, y0, x1, y1], fill=255)
+        inset_x, inset_y = w * 0.275, h * 0.275
+        d.ellipse([x0 + inset_x, y0 + inset_y, x1 - inset_x, y1 - inset_y], fill=0)
+    else:
+        raise ValueError("unknown shape: {}".format(shape))
+    return np.asarray(im) > 0
+
+
+def render_scene(scene: Scene) -> np.ndarray:
+    """Render a Scene to a [64,64,3] uint8 array."""
+    hi = np.empty((_HI, _HI, 3), dtype=np.uint8)
+    for leaf in iter_leaves(scene.tree):
+        x0, y0, x1, y1 = leaf["rect"]
+        hi[y0 * _S : y1 * _S, x0 * _S : x1 * _S] = np.asarray(
+            leaf["fill"], dtype=np.uint8
+        )
+    for obj in scene.objects:
+        mask = _shape_mask(obj["shape"], obj["bbox"])
+        tex = _texture_mask(obj["texture"])
+        rgb = np.asarray(obj["rgb"], dtype=np.float64)
+        main = rgb.astype(np.uint8)
+        dark = (rgb * 0.45).astype(np.uint8)  # texture gaps: darker shade
+        hi[mask & tex] = main
+        hi[mask & ~tex] = dark
+    lo = Image.fromarray(hi).resize((CANVAS, CANVAS), resample=_BOX)
+    return np.asarray(lo, dtype=np.uint8)
+
+
+def generate(
+    global_seed: int,
+    idx: int,
+    tier: Optional[int] = None,
+    tier_mix: Sequence[float] = DEFAULT_TIER_MIX,
+) -> Tuple[Scene, np.ndarray]:
+    """Convenience: sample scene `idx` and render it. Deterministic per idx."""
+    scene = sample_scene(global_seed, idx, tier=tier, tier_mix=tier_mix)
+    return scene, render_scene(scene)

sprig/data/procgen/sampler.py

@@ -0,0 +1,476 @@
+"""Ground-truth BSP scene sampler.
+
+Scenes are generated BY sampling a BSP region tree first (the tree is layout
+sampler, renderer input, and parse-diagnostic label at once), then decorating
+leaves with backgrounds / objects. The tree is constrained to the model
+support of DESIGN.md §3:
+
+- every cut lands on the 8-px grid with relative offset in [0.3, 0.7];
+- leaf regions have both sides in {8, 16} px (any side > 16 must expand,
+  8x8 must terminate);
+- object cells are 16x16 leaves holding exactly one shape, fully inside the
+  cell with >= 2 px margin — no overlap, no occlusion.
+
+Tree JSON schema (serialized verbatim into meta.jsonl):
+  internal node: {"rect": [x0,y0,x1,y1], "axis": "V"|"H", "cut": px,
+                  "children": [lo, hi]}          # lo = left/top child
+  leaf node:     {"rect": [x0,y0,x1,y1], "leaf": true, "obj": int|null,
+                  "fill": [r,g,b]}               # optional "frame": true
+Tier-3 frame region nodes additionally carry {"frame_color": name}.
+
+Object schema: {"shape","color","size","texture": str, "rgb": [r,g,b],
+                "cell": [x0,y0,x1,y1], "bbox": [x0,y0,x1,y1]}.
+"""
+from __future__ import annotations
+
+from dataclasses import dataclass, field
+from typing import Any, Dict, Iterator, List, Optional, Sequence, Tuple
+
+import numpy as np
+
+from .vocab import (
+    BACKGROUNDS,
+    BACKGROUND_NAMES,
+    BG_JITTER,
+    CANVAS,
+    COLOR_JITTER,
+    COLOR_NAMES,
+    COLORS,
+    GRID,
+    HOLDOUT_COMBOS,
+    MAX_LEAF,
+    MIN_MARGIN,
+    OFFSET_HI,
+    OFFSET_LO,
+    SHAPES,
+    SIZE_NAMES,
+    SIZES,
+)
+
+DEFAULT_TIER_MIX: Tuple[float, float, float, float] = (0.10, 0.30, 0.40, 0.20)
+
+# termination probability for leaf-eligible regions in the generic grower
+_TERM_P: Dict[Tuple[int, int], float] = {(16, 16): 0.75, (8, 16): 0.7, (16, 8): 0.7}
+_MAX_TRIES = 200
+_EPS = 1e-9
+
+
+@dataclass
+class Scene:
+    """One procedural scene. All fields are JSON-able plain python."""
+
+    idx: int
+    tier: int
+    tree: Dict[str, Any]
+    objects: List[Dict[str, Any]]
+    background: str  # background name; "sky|sand" for tier-2 sky/ground scenes
+    relation: Optional[Dict[str, Any]] = None  # tier 1: {"type","a","b"}
+    frame: Optional[Dict[str, Any]] = None  # tier 3: {"color","rgb","rect","inner"}
+
+
+# --------------------------------------------------------------------------
+# rng helpers
+# --------------------------------------------------------------------------
+
+def scene_rng(global_seed: int, idx: int) -> np.random.Generator:
+    """The canonical per-sample rng stream: SeedSequence([global_seed, idx])."""
+    return np.random.default_rng(np.random.SeedSequence([int(global_seed), int(idx)]))
+
+
+def caption_rng(global_seed: int, idx: int) -> np.random.Generator:
+    """Independent per-sample stream for caption template draws."""
+    return np.random.default_rng(
+        np.random.SeedSequence([int(global_seed), int(idx), 1])
+    )
+
+
+# --------------------------------------------------------------------------
+# tree construction
+# --------------------------------------------------------------------------
+
+def _mk_leaf(rect: Sequence[int]) -> Dict[str, Any]:
+    return {"rect": [int(v) for v in rect], "leaf": True, "obj": None, "fill": None}
+
+
+def _mk_node(
+    rect: Sequence[int], axis: str, cut: int, lo: Dict[str, Any], hi: Dict[str, Any]
+) -> Dict[str, Any]:
+    return {
+        "rect": [int(v) for v in rect],
+        "axis": axis,
+        "cut": int(cut),
+        "children": [lo, hi],
+    }
+
+
+def _split_rect(
+    rect: Sequence[int], axis: str, cut: int
+) -> Tuple[Tuple[int, int, int, int], Tuple[int, int, int, int]]:
+    x0, y0, x1, y1 = rect
+    if axis == "V":
+        return (x0, y0, cut, y1), (cut, y0, x1, y1)
+    return (x0, y0, x1, cut), (x0, cut, x1, y1)
+
+
+def _valid_cuts(lo: int, hi: int) -> List[int]:
+    """Grid cuts of interval [lo,hi) with relative offset in [0.3, 0.7]."""
+    length = hi - lo
+    out = []
+    for j in range(1, length // GRID):
+        t = j * GRID / length
+        if OFFSET_LO - _EPS <= t <= OFFSET_HI + _EPS:
+            out.append(lo + j * GRID)
+    return out
+
+
+def _grow(
+    rng: np.random.Generator,
+    rect: Tuple[int, int, int, int],
+    force_axis: Optional[str] = None,
+) -> Dict[str, Any]:
+    """Sample a random in-support BSP subtree over `rect`.
+
+    `force_axis` forces the root of this subtree to split on that axis
+    (used for tier-2 sky/ground scenes).
+    """
+    x0, y0, x1, y1 = rect
+    w, h = x1 - x0, y1 - y0
+    leaf_ok = w <= MAX_LEAF and h <= MAX_LEAF
+    if leaf_ok and force_axis is None:
+        if (w == GRID and h == GRID) or rng.random() < _TERM_P[(w, h)]:
+            return _mk_leaf(rect)
+    axes = []
+    if _valid_cuts(x0, x1):
+        axes.append("V")
+    if _valid_cuts(y0, y1):
+        axes.append("H")
+    if force_axis is not None:
+        assert force_axis in axes, "forced axis has no valid cut"
+        axes = [force_axis]
+    axis = axes[int(rng.integers(len(axes)))]
+    cuts = _valid_cuts(x0, x1) if axis == "V" else _valid_cuts(y0, y1)
+    cut = cuts[int(rng.integers(len(cuts)))]
+    rlo, rhi = _split_rect(rect, axis, cut)
+    return _mk_node(rect, axis, cut, _grow(rng, rlo), _grow(rng, rhi))
+
+
+def iter_leaves(node: Dict[str, Any]) -> Iterator[Dict[str, Any]]:
+    """Yield leaf node dicts (mutable references) in depth-first lo-hi order."""
+    if node.get("leaf"):
+        yield node
+    else:
+        for child in node["children"]:
+            for leaf in iter_leaves(child):
+                yield leaf
+
+
+def _object_cells(node: Dict[str, Any]) -> List[Dict[str, Any]]:
+    """Leaves big enough to hold an object (16x16, not frame cells)."""
+    out = []
+    for leaf in iter_leaves(node):
+        x0, y0, x1, y1 = leaf["rect"]
+        if x1 - x0 == MAX_LEAF and y1 - y0 == MAX_LEAF and not leaf.get("frame"):
+            out.append(leaf)
+    return out
+
+
+# --------------------------------------------------------------------------
+# attribute / placement helpers
+# --------------------------------------------------------------------------
+
+def _jitter(rgb: Sequence[int], rng: np.random.Generator, amt: int) -> List[int]:
+    out = []
+    for v in rgb:
+        out.append(int(np.clip(int(v) + int(rng.integers(-amt, amt + 1)), 0, 255)))
+    return out
+
+
+def _sample_attrs(
+    rng: np.random.Generator, exclude: Optional[set] = None
+) -> Dict[str, str]:
+    """Rejection-resample attributes so no holdout (color, shape) combo appears."""
+    exclude = exclude or set()
+    for _ in range(_MAX_TRIES):
+        color = COLOR_NAMES[int(rng.integers(len(COLOR_NAMES)))]
+        shape = SHAPES[int(rng.integers(len(SHAPES)))]
+        if (color, shape) in HOLDOUT_COMBOS or (color, shape) in exclude:
+            continue
+        size = SIZE_NAMES[int(rng.integers(len(SIZE_NAMES)))]
+        texture = "solid" if rng.random() < 0.4 else (
+            ("striped", "dotted", "checker")[int(rng.integers(3))]
+        )
+        return {"shape": shape, "color": color, "size": size, "texture": texture}
+    raise RuntimeError("attribute rejection loop exhausted")
+
+
+def _place_object(
+    rng: np.random.Generator,
+    cell: Sequence[int],
+    attrs: Dict[str, str],
+    centered: bool = False,
+) -> Dict[str, Any]:
+    x0, y0, x1, y1 = (float(v) for v in cell)
+    s = SIZES[attrs["size"]]
+    w = s
+    h = round(s * 0.62, 2) if attrs["shape"] == "rectangle" else s
+    slack_x = (x1 - x0 - w) / 2.0 - MIN_MARGIN
+    slack_y = (y1 - y0 - h) / 2.0 - MIN_MARGIN
+    assert slack_x >= -_EPS and slack_y >= -_EPS, "object does not fit cell"
+    cx = (x0 + x1) / 2.0
+    cy = (y0 + y1) / 2.0
+    if not centered:
+        cx += float(rng.uniform(-max(slack_x, 0.0), max(slack_x, 0.0)))
+        cy += float(rng.uniform(-max(slack_y, 0.0), max(slack_y, 0.0)))
+    bbox = [
+        round(cx - w / 2.0, 2),
+        round(cy - h / 2.0, 2),
+        round(cx + w / 2.0, 2),
+        round(cy + h / 2.0, 2),
+    ]
+    return {
+        "shape": attrs["shape"],
+        "color": attrs["color"],
+        "size": attrs["size"],
+        "texture": attrs["texture"],
+        "rgb": _jitter(COLORS[attrs["color"]], rng, COLOR_JITTER),
+        "cell": [int(v) for v in cell],
+        "bbox": bbox,
+    }
+
+
+def _apply_bg(tree: Dict[str, Any], bg_rgb: Sequence[int]) -> None:
+    for leaf in iter_leaves(tree):
+        if leaf["fill"] is None:
+            leaf["fill"] = list(bg_rgb)
+
+
+def _pick_bg(rng: np.random.Generator) -> Tuple[str, List[int]]:
+    name = BACKGROUND_NAMES[int(rng.integers(len(BACKGROUND_NAMES)))]
+    return name, _jitter(BACKGROUNDS[name], rng, BG_JITTER)
+
+
+# --------------------------------------------------------------------------
+# tier builders
+# --------------------------------------------------------------------------
+
+def _build_tier0(rng: np.random.Generator, idx: int) -> Scene:
+    for _ in range(_MAX_TRIES):
+        tree = _grow(rng, (0, 0, CANVAS, CANVAS))
+        cells = _object_cells(tree)
+        if cells:
+            break
+    else:  # pragma: no cover
+        raise RuntimeError("tier0: no object cell after retries")
+    c = CANVAS / 2.0
+    cell = min(
+        cells,
+        key=lambda l: ((l["rect"][0] + l["rect"][2]) / 2 - c) ** 2
+        + ((l["rect"][1] + l["rect"][3]) / 2 - c) ** 2,
+    )
+    obj = _place_object(rng, cell["rect"], _sample_attrs(rng), centered=True)
+    cell["obj"] = 0
+    bg_name, bg_rgb = _pick_bg(rng)
+    _apply_bg(tree, bg_rgb)
+    return Scene(idx=idx, tier=0, tree=tree, objects=[obj], background=bg_name)
+
+
+def _build_tier1(rng: np.random.Generator, idx: int) -> Scene:
+    for _ in range(_MAX_TRIES):
+        tree = _grow(rng, (0, 0, CANVAS, CANVAS))
+        cells_lo = _object_cells(tree["children"][0])
+        cells_hi = _object_cells(tree["children"][1])
+        if cells_lo and cells_hi:
+            break
+    else:  # pragma: no cover
+        raise RuntimeError("tier1: no cells on both sides of root split")
+    cell_a = cells_lo[int(rng.integers(len(cells_lo)))]
+    cell_b = cells_hi[int(rng.integers(len(cells_hi)))]
+    attrs_a = _sample_attrs(rng)
+    attrs_b = _sample_attrs(rng, exclude={(attrs_a["color"], attrs_a["shape"])})
+    obj_a = _place_object(rng, cell_a["rect"], attrs_a)
+    obj_b = _place_object(rng, cell_b["rect"], attrs_b)
+    cell_a["obj"] = 0
+    cell_b["obj"] = 1
+    relation = {"type": "left" if tree["axis"] == "V" else "above", "a": 0, "b": 1}
+    bg_name, bg_rgb = _pick_bg(rng)
+    _apply_bg(tree, bg_rgb)
+    return Scene(
+        idx=idx,
+        tier=1,
+        tree=tree,
+        objects=[obj_a, obj_b],
+        background=bg_name,
+        relation=relation,
+    )
+
+
+def _build_tier2(rng: np.random.Generator, idx: int) -> Scene:
+    n_obj = 3 + int(rng.integers(3))  # 3..5
+    skyground = rng.random() < 0.5
+    for _ in range(_MAX_TRIES):
+        tree = _grow(
+            rng, (0, 0, CANVAS, CANVAS), force_axis="H" if skyground else None
+        )
+        cells = _object_cells(tree)
+        if len(cells) >= n_obj:
+            break
+    else:  # pragma: no cover
+        raise RuntimeError("tier2: not enough object cells")
+    order = rng.permutation(len(cells))[:n_obj]
+    objects: List[Dict[str, Any]] = []
+    used: set = set()
+    for i, ci in enumerate(order):
+        attrs = _sample_attrs(rng, exclude=used)
+        used.add((attrs["color"], attrs["shape"]))
+        cell = cells[int(ci)]
+        objects.append(_place_object(rng, cell["rect"], attrs))
+        cell["obj"] = i
+    if skyground:
+        cut = tree["cut"]
+        sky_rgb = _jitter(BACKGROUNDS["sky"], rng, BG_JITTER)
+        sand_rgb = _jitter(BACKGROUNDS["sand"], rng, BG_JITTER)
+        for leaf in iter_leaves(tree):
+            leaf["fill"] = list(sky_rgb) if leaf["rect"][3] <= cut else list(sand_rgb)
+        bg_name = "sky|sand"
+    else:
+        bg_name, bg_rgb = _pick_bg(rng)
+        _apply_bg(tree, bg_rgb)
+    return Scene(idx=idx, tier=2, tree=tree, objects=objects, background=bg_name)
+
+
+def _frame_region_tree(
+    rng: np.random.Generator,
+    rect: Tuple[int, int, int, int],
+    frame_rgb: List[int],
+) -> Tuple[Dict[str, Any], Dict[str, Any]]:
+    """Decompose a 40x40 region into frame cells around one inner 16x16 cell.
+
+    Per axis the strips are (16, 16, 8) or (8, 16, 16) so every cut keeps its
+    relative offset in [0.3, 0.7]. Returns (subtree, inner_leaf_ref).
+    """
+    fx, fy, _, _ = rect
+
+    def frame_leaf(r: Sequence[int]) -> Dict[str, Any]:
+        leaf = _mk_leaf(r)
+        leaf["frame"] = True
+        leaf["fill"] = list(frame_rgb)
+        return leaf
+
+    def frame_subtree(r: Tuple[int, int, int, int]) -> Dict[str, Any]:
+        sub = _grow(rng, r)
+        for leaf in iter_leaves(sub):
+            leaf["frame"] = True
+            leaf["fill"] = list(frame_rgb)
+        return sub
+
+    f0x = 16 if rng.random() < 0.5 else 8
+    f0y = 16 if rng.random() < 0.5 else 8
+    ix0 = fx + f0x  # inner cell x range [ix0, ix0+16]
+    iy0 = fy + f0y
+
+    # middle vertical strip (16 wide, 40 tall) split along y into
+    # top frame cell / inner object cell / bottom frame cell
+    inner = _mk_leaf((ix0, iy0, ix0 + 16, iy0 + 16))
+    top = frame_leaf((ix0, fy, ix0 + 16, iy0))
+    bot = frame_leaf((ix0, iy0 + 16, ix0 + 16, fy + 40))
+    strip_rect = (ix0, fy, ix0 + 16, fy + 40)
+    if f0y == 16:
+        lower = _mk_node((ix0, iy0, ix0 + 16, fy + 40), "H", iy0 + 16, inner, bot)
+        mid = _mk_node(strip_rect, "H", iy0, top, lower)
+    else:
+        upper = _mk_node((ix0, fy, ix0 + 16, iy0 + 16), "H", iy0, top, inner)
+        mid = _mk_node(strip_rect, "H", iy0 + 16, upper, bot)
+
+    left = frame_subtree((fx, fy, ix0, fy + 40))
+    right = frame_subtree((ix0 + 16, fy, fx + 40, fy + 40))
+    if f0x == 16:
+        rt = _mk_node((ix0, fy, fx + 40, fy + 40), "V", ix0 + 16, mid, right)
+        node = _mk_node(rect, "V", ix0, left, rt)
+    else:
+        lt = _mk_node((fx, fy, ix0 + 16, fy + 40), "V", ix0, left, mid)
+        node = _mk_node(rect, "V", ix0 + 16, lt, right)
+    return node, inner
+
+
+def _build_tier3(rng: np.random.Generator, idx: int) -> Scene:
+    fx = 0 if rng.random() < 0.5 else 24
+    fy = 0 if rng.random() < 0.5 else 24
+    frect = (fx, fy, fx + 40, fy + 40)
+
+    attrs = _sample_attrs(rng)
+    frame_color = attrs["color"]
+    while frame_color == attrs["color"]:
+        frame_color = COLOR_NAMES[int(rng.integers(len(COLOR_NAMES)))]
+    frame_rgb = _jitter(COLORS[frame_color], rng, COLOR_JITTER)
+
+    frame_node, inner = _frame_region_tree(rng, frect, frame_rgb)
+    frame_node["frame_color"] = frame_color
+    obj = _place_object(rng, inner["rect"], attrs)
+    inner["obj"] = 0
+
+    # isolate the frame region from the canvas root: one V cut, one H cut
+    # (in random order); both offsets (24/64, 40/64) are in [0.3, 0.7].
+    def isolate(rect: Tuple[int, int, int, int], axis: str, inner_node: Dict) -> Dict:
+        x0, y0, x1, y1 = rect
+        if axis == "V":
+            cut = fx + 40 if fx == 0 else fx
+            (rl, rh) = _split_rect(rect, "V", cut)
+            keep_lo = fx == 0
+        else:
+            cut = fy + 40 if fy == 0 else fy
+            (rl, rh) = _split_rect(rect, "H", cut)
+            keep_lo = fy == 0
+        other = _grow(rng, rh if keep_lo else rl)
+        lo = inner_node if keep_lo else other
+        hi = other if keep_lo else inner_node
+        return _mk_node(rect, axis, cut, lo, hi)
+
+    if rng.random() < 0.5:
+        band = isolate((fx, 0, fx + 40, CANVAS), "H", frame_node)
+        tree = isolate((0, 0, CANVAS, CANVAS), "V", band)
+    else:
+        band = isolate((0, fy, CANVAS, fy + 40), "V", frame_node)
+        tree = isolate((0, 0, CANVAS, CANVAS), "H", band)
+
+    bg_name, bg_rgb = _pick_bg(rng)
+    _apply_bg(tree, bg_rgb)
+    frame_info = {
+        "color": frame_color,
+        "rgb": list(frame_rgb),
+        "rect": list(frect),
+        "inner": list(inner["rect"]),
+    }
+    return Scene(
+        idx=idx,
+        tier=3,
+        tree=tree,
+        objects=[obj],
+        background=bg_name,
+        frame=frame_info,
+    )
+
+
+_BUILDERS = {0: _build_tier0, 1: _build_tier1, 2: _build_tier2, 3: _build_tier3}
+
+
+# --------------------------------------------------------------------------
+# public API
+# --------------------------------------------------------------------------
+
+def sample_scene(
+    global_seed: int,
+    idx: int,
+    tier: Optional[int] = None,
+    tier_mix: Sequence[float] = DEFAULT_TIER_MIX,
+) -> Scene:
+    """Deterministically sample scene `idx` under `global_seed`.
+
+    If `tier` is None it is drawn from `tier_mix` (first rng draw, so the
+    scene is a pure function of (global_seed, idx, tier or tier_mix)).
+    """
+    rng = scene_rng(global_seed, idx)
+    if tier is None:
+        p = np.asarray(tier_mix, dtype=np.float64)
+        p = p / p.sum()
+        tier = int(rng.choice(4, p=p))
+    return _BUILDERS[int(tier)](rng, int(idx))

sprig/data/procgen/vocab.py

@@ -0,0 +1,72 @@
+"""Vocabulary and geometry constants for the procedural 2D scene generator.
+
+Everything here is a plain constant so the sampler / renderer / captioner and
+the eval agents (color anchor classification) share one source of truth.
+"""
+from __future__ import annotations
+
+from typing import Dict, Tuple
+
+# --- geometry (must match the model lattice in DESIGN.md §2/§3) -------------
+CANVAS: int = 64          # canvas side, px
+GRID: int = 8             # model lattice stride, px — ALL cuts land on this grid
+MAX_LEAF: int = 16        # model support: leaf regions have both sides <= 16 px
+SUPERSAMPLE: int = 4      # render at 256x256, BOX-downsample to 64x64
+OFFSET_LO: float = 0.3    # relative cut offset range (inside model support)
+OFFSET_HI: float = 0.7
+MIN_MARGIN: float = 2.0   # min px margin between a shape bbox and its cell walls
+
+# --- object vocabulary -------------------------------------------------------
+# 8 colors: name -> RGB anchor (renderer jitters around these; eval classifies
+# back to nearest anchor in Lab space).
+COLORS: Dict[str, Tuple[int, int, int]] = {
+    "red": (210, 45, 45),
+    "green": (55, 170, 60),
+    "blue": (50, 90, 215),
+    "yellow": (235, 215, 55),
+    "orange": (240, 145, 40),
+    "purple": (140, 60, 185),
+    "cyan": (60, 205, 215),
+    "magenta": (220, 70, 175),
+}
+COLOR_NAMES: Tuple[str, ...] = tuple(COLORS.keys())
+
+SHAPES: Tuple[str, ...] = (
+    "circle", "square", "triangle", "rectangle", "diamond", "star", "cross", "ring",
+)
+
+# 2 sizes: name -> shape bbox side in px. Objects live in 16x16 leaf cells with
+# >= MIN_MARGIN px of clearance: large 11 -> 2.5 px margin, small 6 -> 5 px.
+SIZES: Dict[str, float] = {"small": 6.0, "large": 11.0}
+SIZE_NAMES: Tuple[str, ...] = tuple(SIZES.keys())
+
+TEXTURES: Tuple[str, ...] = ("solid", "striped", "dotted", "checker")
+# caption adjective per texture ("" = no adjective for solid fills)
+TEXTURE_ADJ: Dict[str, str] = {
+    "solid": "",
+    "striped": "striped",
+    "dotted": "dotted",
+    "checker": "checkered",
+}
+
+# 5 backgrounds: name -> RGB anchor. "sky"/"sand" are also used as the fixed
+# pair for the tier-2 sky/ground root split.
+BACKGROUNDS: Dict[str, Tuple[int, int, int]] = {
+    "white": (245, 245, 245),
+    "black": (28, 28, 28),
+    "gray": (128, 128, 128),
+    "sky": (170, 210, 240),
+    "sand": (205, 180, 130),
+}
+BACKGROUND_NAMES: Tuple[str, ...] = tuple(BACKGROUNDS.keys())
+
+# per-channel uniform jitter amplitude applied to anchors at scene-sample time
+COLOR_JITTER: int = 12
+BG_JITTER: int = 8
+
+# --- compositional holdout ---------------------------------------------------
+# These (color, shape) combos NEVER appear in training scenes or captions
+# (enforced by rejection-resampling in the sampler; scanned in tests/CI).
+HOLDOUT_COMBOS: frozenset = frozenset(
+    {("blue", "triangle"), ("red", "ring"), ("green", "star"), ("yellow", "cross")}
+)

sprig/data/procgen/writer.py

@@ -0,0 +1,193 @@
+"""Parallel pregeneration of the procedural dataset to raw memmaps.
+
+Outputs in --out DIR:
+  images.u8            uint8 memmap [N,64,64,3]
+  meta.jsonl           one JSON object per line:
+                       {idx, tier, caption, template_id, partial, objects, tree}
+  meta_offsets.i64     int64 [N]: byte offset of the START of line i
+  tier_idx/tier{0..3}.i64  int64 sorted idx arrays per tier (curriculum)
+
+Embeddings are NOT written here — sprig/data/embed_t5.py does that later.
+
+CLI:
+  python -m sprig.data.procgen.writer --out DIR --n N --seed SEED \
+      --tier-mix "0.1,0.3,0.4,0.2" --workers K
+
+Workers own disjoint contiguous row ranges; each writes its rows of the image
+memmap plus a meta/tier shard, and the parent concatenates shards in order.
+Every sample is a pure function of (seed, idx), so the output is independent
+of the worker count.
+"""
+from __future__ import annotations
+
+import argparse
+import json
+import multiprocessing as mp
+import os
+import sys
+import time
+from typing import List, Optional, Sequence, Tuple
+
+import numpy as np
+
+from .captions import sample_caption
+from .render import render_scene
+from .sampler import DEFAULT_TIER_MIX, caption_rng, sample_scene
+
+IMG_SHAPE = (64, 64, 3)
+
+
+def _meta_part(out_dir: str, wid: int) -> str:
+    return os.path.join(out_dir, "meta.part{:03d}.jsonl".format(wid))
+
+
+def _tier_part(out_dir: str, wid: int) -> str:
+    return os.path.join(out_dir, "tiers.part{:03d}.i64".format(wid))
+
+
+def _generate_range(
+    out_dir: str,
+    n: int,
+    seed: int,
+    tier_mix: Sequence[float],
+    lo: int,
+    hi: int,
+    wid: int,
+    seed_start: int = 0,
+) -> None:
+    """Worker: generate rows [lo, hi) into the shared image memmap + shards.
+
+    Row idx is generated from scene index seed_start + idx, so different
+    splits get disjoint scene streams from disjoint seed ranges."""
+    images = np.memmap(
+        os.path.join(out_dir, "images.u8"), dtype=np.uint8, mode="r+",
+        shape=(n,) + IMG_SHAPE,
+    )
+    tiers = np.empty(hi - lo, dtype=np.int64)
+    with open(_meta_part(out_dir, wid), "w", encoding="utf-8") as f:
+        for idx in range(lo, hi):
+            sidx = seed_start + idx
+            scene = sample_scene(seed, sidx, tier_mix=tier_mix)
+            images[idx] = render_scene(scene)
+            cap = sample_caption(scene, caption_rng(seed, sidx), mode="train")
+            rec = {
+                "idx": idx,
+                "tier": scene.tier,
+                "caption": cap.text,
+                "template_id": cap.template_id,
+                "partial": cap.partial,
+                "objects": scene.objects,
+                "tree": scene.tree,
+            }
+            f.write(json.dumps(rec, separators=(",", ":")) + "\n")
+            tiers[idx - lo] = scene.tier
+    images.flush()
+    del images
+    tiers.tofile(_tier_part(out_dir, wid))
+
+
+def write_dataset(
+    out_dir: str,
+    n: int,
+    seed: int = 0,
+    tier_mix: Sequence[float] = DEFAULT_TIER_MIX,
+    workers: int = 1,
+    seed_start: int = 0,
+) -> None:
+    os.makedirs(out_dir, exist_ok=True)
+    os.makedirs(os.path.join(out_dir, "tier_idx"), exist_ok=True)
+    workers = max(1, min(int(workers), n))
+
+    # preallocate the image memmap
+    images = np.memmap(
+        os.path.join(out_dir, "images.u8"), dtype=np.uint8, mode="w+",
+        shape=(n,) + IMG_SHAPE,
+    )
+    images.flush()
+    del images
+
+    # disjoint contiguous row ranges
+    bounds = np.linspace(0, n, workers + 1).astype(np.int64)
+    ranges: List[Tuple[int, int]] = [
+        (int(bounds[w]), int(bounds[w + 1])) for w in range(workers)
+    ]
+
+    t0 = time.time()
+    if workers == 1:
+        _generate_range(out_dir, n, seed, tier_mix, 0, n, 0, seed_start)
+    else:
+        ctx = mp.get_context("spawn")
+        procs = []
+        for wid, (lo, hi) in enumerate(ranges):
+            p = ctx.Process(
+                target=_generate_range,
+                args=(out_dir, n, seed, tuple(tier_mix), lo, hi, wid, seed_start),
+            )
+            p.start()
+            procs.append(p)
+        for p in procs:
+            p.join()
+        failed = [p.exitcode for p in procs if p.exitcode != 0]
+        if failed:
+            raise RuntimeError("worker(s) failed with exit codes {}".format(failed))
+
+    # concatenate meta shards in order, recording line-start byte offsets
+    offsets = np.empty(n, dtype=np.int64)
+    tiers = np.empty(n, dtype=np.int64)
+    pos = 0
+    row = 0
+    with open(os.path.join(out_dir, "meta.jsonl"), "wb") as out:
+        for wid in range(workers):
+            with open(_meta_part(out_dir, wid), "rb") as part:
+                for line in part:
+                    offsets[row] = pos
+                    out.write(line)
+                    pos += len(line)
+                    row += 1
+            tiers_part = np.fromfile(_tier_part(out_dir, wid), dtype=np.int64)
+            lo, hi = ranges[wid]
+            tiers[lo:hi] = tiers_part
+            os.remove(_meta_part(out_dir, wid))
+            os.remove(_tier_part(out_dir, wid))
+    assert row == n, "meta line count {} != n {}".format(row, n)
+    offsets.tofile(os.path.join(out_dir, "meta_offsets.i64"))
+    for t in range(4):
+        idxs = np.nonzero(tiers == t)[0].astype(np.int64)
+        idxs.tofile(os.path.join(out_dir, "tier_idx", "tier{}.i64".format(t)))
+
+    counts = [int((tiers == t).sum()) for t in range(4)]
+    print(
+        "wrote {} samples to {} in {:.1f}s (tier counts: {})".format(
+            n, out_dir, time.time() - t0, counts
+        )
+    )
+
+
+def _parse_tier_mix(s: str) -> Tuple[float, ...]:
+    parts = tuple(float(v) for v in s.split(","))
+    if len(parts) != 4 or min(parts) < 0 or sum(parts) <= 0:
+        raise argparse.ArgumentTypeError(
+            "--tier-mix must be 4 nonnegative comma-separated floats"
+        )
+    return parts
+
+
+def main(argv: Optional[Sequence[str]] = None) -> None:
+    ap = argparse.ArgumentParser(description="Pregenerate the proc2d dataset.")
+    ap.add_argument("--out", required=True, help="output directory")
+    ap.add_argument("--n", type=int, required=True, help="number of samples")
+    ap.add_argument("--seed", type=int, default=0)
+    ap.add_argument(
+        "--tier-mix", type=_parse_tier_mix,
+        default=DEFAULT_TIER_MIX, help='e.g. "0.1,0.3,0.4,0.2"',
+    )
+    ap.add_argument("--workers", type=int, default=1)
+    ap.add_argument("--seed-start", type=int, default=0,
+                    help="scene-index offset (disjoint ranges per split)")
+    args = ap.parse_args(argv)
+    write_dataset(args.out, args.n, args.seed, args.tier_mix, args.workers,
+                  seed_start=args.seed_start)
+
+
+if __name__ == "__main__":
+    main(sys.argv[1:])

sprig/dp/inside.py

@@ -0,0 +1,604 @@
+"""Inside DP over the finite region lattice (DESIGN.md section 5).
+
+Exposes the three DP entry points SPRIGModel calls through its adapter
+(sprig.model.sprig._call_dp), all keyed by the DESIGN section 5 argument
+names plus ``log_PT`` (the [R, T_v] log texel prior table the leaf T-mix
+needs; p(T|A,c) = logsumexp_k(U_logmix[..,A,k] + log_PT[k,T])):
+
+    inside(ell_leaf, term_logits, cut_logits, U_logmix, logV, logW,
+           lattice, temper_kappa, log_PT)      -> (beta [B,N_reg,S] fp32, logZ [B])
+    inside_logZ(...)                           -> alias of ``inside``
+    viterbi(...)                               -> (score [B], trees)
+    posterior_marginals(...)                   -> dict of expected-count tensors
+
+Semantics (matching tests/fixtures_dp_stub.py, the brute-force-validated
+reference):
+  * beta(r, A) for leaf-eligible r combines
+        log p_term(r,A) + logsumexp_T(log p(T|A,c) + ell(r,T)/kappa(r))
+    with
+        log(1 - p_term(r,A)) + expand(r,A),
+    where must_terminate regions force p_term = 1 and must_expand regions
+    (any side > leaf_max) expand with probability 1 (no p_term factor).
+  * expand(r, A) is computed by a level-synchronous sweep in cell-area
+    ascending order using the lattice's flattened per-level cut tables;
+    cut-type probability mass is split uniformly across same-type concrete
+    cuts (the precomputed -log(count) correction).
+  * logZ[b] = beta[b, root, 0] (axiom nonterminal = symbol 0).
+
+Numerics: all log-domain accumulation in fp32; every logsumexp/logbmm is
+max-shifted with *detached* shifts. -1e30 sentinels are never fed through
+branches on the differentiable path (they only seed the beta table and the
+scatter-max identity); leaf/expand branches are combined by mask + index so
+the module is safe under double backward (SPRIGModel.loss uses
+create_graph=True through the DP for the under-use hinges).
+
+Every function accepts ``exact=True`` to replace the max-shifted
+exp/matmul/log fast path with a direct fp32 logsumexp — the exact fallback
+DESIGN section 5 asks for (used by the parity tests).
+"""
+from __future__ import annotations
+
+from typing import Dict, List, Optional, Tuple
+
+import torch
+import torch.nn.functional as F
+
+NEG = -1e30
+# Clamp floor for post-shift exp-sums before the final log. It must be large
+# enough that log's DOUBLE backward (-grad/x^2) stays finite in fp32
+# (x^2 >= 1e-30 > fp32 tiny), because SPRIGModel.loss differentiates the
+# under-use hinges through the DP with create_graph=True. Floor 1e-15 caps a
+# reduction at ~34.5 nats below the factored max-shifts, which only rounds up
+# paths of negligible posterior mass. (The segment-logsumexp sums are >= 1 by
+# construction and never hit the clamp.)
+_TINY = 1e-15
+# Element budget for the [B, chunk, S, R] / [B, chunk, S, T_v] blocks the
+# max-semiring (argmax) passes materialize.
+_MAX_SEMIRING_BUDGET = 1 << 24
+
+
+def logbmm(x_log: torch.Tensor, w_log: torch.Tensor, exact: bool = False) -> torch.Tensor:
+    """out[..., m] = logsumexp_k(x_log[..., k] + w_log[k, m]).
+
+    x_log [..., K] and w_log [K, M], both finite fp32. Max-shifted (detached
+    shifts), exp, matmul, log, shifts added back. ``exact=True`` is the direct
+    fp32 logsumexp fallback for tests.
+    """
+    x_log = x_log.float()
+    w_log = w_log.float()
+    if exact:
+        expand_shape = (1,) * (x_log.dim() - 1) + tuple(w_log.shape)
+        return torch.logsumexp(x_log.unsqueeze(-1) + w_log.view(expand_shape), dim=-2)
+    sx = x_log.max(dim=-1, keepdim=True).values.detach()          # [..., 1]
+    sw = w_log.max(dim=0, keepdim=True).values.detach()           # [1, M]
+    out = torch.matmul(torch.exp(x_log - sx), torch.exp(w_log - sw))
+    return torch.log(out.clamp(min=_TINY)) + sx + sw
+
+
+def logbmm_batched(x_log: torch.Tensor, w_log: torch.Tensor, exact: bool = False) -> torch.Tensor:
+    """out[b, p, m] = logsumexp_k(x_log[b, p, k] + w_log[b, k, m]).
+
+    Batched-weight variant of :func:`logbmm` (x [B, P, K], w [B, K, M])."""
+    x_log = x_log.float()
+    w_log = w_log.float()
+    if exact:
+        return torch.logsumexp(x_log.unsqueeze(-1) + w_log.unsqueeze(1), dim=-2)
+    sx = x_log.max(dim=-1, keepdim=True).values.detach()          # [B, P, 1]
+    sw = w_log.max(dim=1, keepdim=True).values.detach()           # [B, 1, M]
+    out = torch.bmm(torch.exp(x_log - sx), torch.exp(w_log - sw))
+    return torch.log(out.clamp(min=_TINY)) + sx + sw
+
+
+def _segment_logsumexp(
+    x: torch.Tensor, seg_index: torch.Tensor, n_seg: int
+) -> torch.Tensor:
+    """Group-by-segment logsumexp along dim 1: x [B, M, R], seg_index [M]
+    -> [B, n_seg, R]. Max-shift per segment (detached), scatter-add of exp,
+    log."""
+    B, M, R = x.shape
+    idx = seg_index.view(1, -1, 1).expand(B, M, R)
+    mx = torch.full((B, n_seg, R), NEG, device=x.device, dtype=x.dtype).scatter_reduce(
+        1, idx, x.detach(), reduce="amax", include_self=True
+    )
+    ex = torch.exp(x - mx[:, seg_index, :])
+    ssum = torch.zeros(B, n_seg, R, device=x.device, dtype=x.dtype).index_add_(
+        1, seg_index, ex
+    )
+    return mx + torch.log(ssum.clamp(min=_TINY))
+
+
+def _row_cut_logp(cut_logits: torch.Tensor, lattice, lv, exact: bool = False) -> torch.Tensor:
+    """Per concrete cut row: masked + renormalized cut-type log-prob with the
+    uniform same-type mass split -> [B, M, R].
+
+    cut_logits [B, R, 14]; validity mask per unique parent from
+    lattice.type_present; log-softmax over the masked set; -log(count)
+    correction per row (lv.log_cnt). (Fallback path — the sweep normally uses
+    :func:`_cut_logp_flat` + one gather per level.)"""
+    del exact  # the masked log-softmax is already exact fp32
+    B, R = cut_logits.shape[0], cut_logits.shape[1]
+    M = lv.parent_ids.shape[0]
+    pres = lattice.type_present[lv.parents]                       # [P, 14]
+    tl = cut_logits.float().unsqueeze(1).masked_fill(
+        ~pres.view(1, -1, 1, pres.shape[-1]), float("-inf")
+    )                                                             # [B, P, R, 14]
+    tls = F.log_softmax(tl, dim=-1)
+    sel = tls[:, lv.parent_index, :, :]                           # [B, M, R, 14]
+    idx = lv.cut_type.view(1, -1, 1, 1).expand(B, M, R, 1)
+    return sel.gather(-1, idx).squeeze(-1) - lv.log_cnt.view(1, -1, 1)
+
+
+def _cut_logp_flat(cut_logits: torch.Tensor, lattice) -> Optional[torch.Tensor]:
+    """Masked + renormalized cut-type log-probs for every DISTINCT
+    type-presence pattern, flattened for per-level gathering:
+
+        out [B, n_pat * 14, R];  row(pattern p, type t) = out[:, p*14 + t, :].
+
+    One log-softmax per sweep instead of one per level; identical values (the
+    per-parent mask in :func:`_row_cut_logp` only depends on the parent's
+    type-presence pattern). Returns None when the lattice does not carry the
+    precomputed pattern tables (foreign/stub lattices -> per-level fallback).
+    """
+    pats = getattr(lattice, "type_patterns", None)
+    if pats is None:
+        return None
+    B, R = cut_logits.shape[0], cut_logits.shape[1]
+    n_pat, T = pats.shape
+    tl = cut_logits.float().unsqueeze(1).masked_fill(
+        ~pats.view(1, n_pat, 1, T), float("-inf")
+    )                                                             # [B, n_pat, R, 14]
+    tls = F.log_softmax(tl, dim=-1)
+    return tls.permute(0, 1, 3, 2).reshape(B, n_pat * T, R)
+
+
+def _level_cut_logp(
+    tls_flat: Optional[torch.Tensor], cut_logits: torch.Tensor, lattice, lv,
+) -> torch.Tensor:
+    """Per-row cut log-prob [B, M, R] for one level (gather from the per-sweep
+    flat table when available, else the per-level fallback)."""
+    if tls_flat is not None and getattr(lv, "flat_type_idx", None) is not None:
+        return tls_flat[:, lv.flat_type_idx, :] - lv.log_cnt.view(1, -1, 1)
+    return _row_cut_logp(cut_logits, lattice, lv)
+
+
+def _mt_meta(lattice) -> Tuple[torch.Tensor, torch.Tensor]:
+    """(must-terminate region ids, their leaf slots) — precomputed on the
+    lattice when available (no per-call ``nonzero``)."""
+    rid_mt = getattr(lattice, "mt_ids", None)
+    if rid_mt is None:
+        rid_mt = torch.nonzero(lattice.must_terminate, as_tuple=False).reshape(-1)
+    slot_mt = getattr(lattice, "mt_slots", None)
+    if slot_mt is None:
+        slot_mt = lattice.leaf_index_of_region[rid_mt]
+    return rid_mt, slot_mt
+
+
+def _sweep_meta(lattice) -> List[Tuple]:
+    """Per-level (has_leaf, has_nonleaf, leaf_pos, nonleaf_pos, leaf_rids,
+    nonleaf_rids, leaf_slots). Reads the tensors precomputed at lattice build;
+    falls back to computing them (with syncs) for foreign/stub lattices."""
+    metas: List[Tuple] = []
+    for lv in lattice.levels:
+        if getattr(lv, "leaf_pos", None) is not None:
+            metas.append((lv.has_leaf, lv.has_nonleaf, lv.leaf_pos,
+                          lv.nonleaf_pos, lv.leaf_rids, lv.nonleaf_rids,
+                          lv.leaf_slots))
+        else:
+            is_leaf = lattice.leaf_mask[lv.parents]
+            leaf_pos = torch.nonzero(is_leaf, as_tuple=False).reshape(-1)
+            nonleaf_pos = torch.nonzero(~is_leaf, as_tuple=False).reshape(-1)
+            leaf_rids = lv.parents[leaf_pos]
+            nonleaf_rids = lv.parents[nonleaf_pos]
+            leaf_slots = lattice.leaf_index_of_region[leaf_rids]
+            metas.append((bool(leaf_pos.numel() > 0), bool(nonleaf_pos.numel() > 0),
+                          leaf_pos, nonleaf_pos, leaf_rids, nonleaf_rids, leaf_slots))
+    return metas
+
+
+def _leaf_terms(
+    ell_leaf: torch.Tensor,
+    term_logits: torch.Tensor,
+    U_logmix: torch.Tensor,
+    log_PT: torch.Tensor,
+    lattice,
+    temper_kappa: torch.Tensor,
+    term_mark: Optional[torch.Tensor] = None,
+    semiring: str = "sum",
+    exact: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor, Optional[torch.Tensor]]:
+    """Termination scores for every leaf-eligible region.
+
+    Returns (term_score [B, n_leaf, S], log_cont [B, n_leaf, S], argT).
+    term_score = log p_term + mix_T(log p(T|A,c) + ell/kappa) where mix is
+    logsumexp (sum semiring) or max (max semiring; argT returned).
+    must_terminate: log p_term = 0, log(1 - p_term) = NEG sentinel (never fed
+    through logaddexp — the caller indexes around it).
+    """
+    prior = logbmm(U_logmix.float(), log_PT.float(), exact=exact)  # [B, S, T_v]
+    ell_t = ell_leaf.float() / temper_kappa.float().view(1, -1, 1)  # [B, n_leaf, T_v]
+    if semiring == "sum":
+        mix = logbmm_batched(ell_t, prior.transpose(1, 2), exact=exact)  # [B, n_leaf, S]
+        argT = None
+    else:
+        B, n_leaf, T_v = ell_t.shape
+        S = prior.shape[1]
+        step = max(1, _MAX_SEMIRING_BUDGET // max(1, B * S * T_v))
+        mixes, args = [], []
+        for i0 in range(0, n_leaf, step):
+            scores = prior.unsqueeze(1) + ell_t[:, i0 : i0 + step].unsqueeze(2)
+            m, a = scores.max(dim=-1)                             # [B, chunk, S]
+            mixes.append(m)
+            args.append(a)
+        mix = torch.cat(mixes, dim=1)
+        argT = torch.cat(args, dim=1)
+    lt = term_logits.float()[:, lattice.leaf_ids, :]
+    log_term = F.logsigmoid(lt)
+    log_cont = F.logsigmoid(-lt)
+    mt = lattice.must_terminate[lattice.leaf_ids].view(1, -1, 1)
+    log_term = torch.where(mt, torch.zeros_like(log_term), log_term)
+    log_cont = torch.where(mt, torch.full_like(log_cont, NEG), log_cont)
+    term_score = log_term + mix
+    if term_mark is not None:
+        term_score = term_score + term_mark[:, lattice.leaf_ids, :]
+    return term_score, log_cont, argT
+
+
+def _sweep(
+    ell_leaf: torch.Tensor,
+    term_logits: torch.Tensor,
+    cut_logits: torch.Tensor,
+    U_logmix: torch.Tensor,
+    logV: torch.Tensor,
+    logW: torch.Tensor,
+    lattice,
+    temper_kappa: torch.Tensor,
+    log_PT: torch.Tensor,
+    term_mark: Optional[torch.Tensor] = None,
+    expand_mark: Optional[torch.Tensor] = None,
+    cut_mark: Optional[torch.Tensor] = None,
+    exact: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Level-synchronous inside sweep -> (beta [B, N_reg, S] fp32, logZ [B]).
+
+    All shape-dependent decisions (leaf/non-leaf membership per level,
+    must-terminate ids, cut-type masks) come from lattice-precomputed index
+    tensors and python bools so the loop issues no CPU-GPU syncs (no
+    ``bool(tensor)``, no bool-mask advanced indexing, no ``nonzero``).
+
+    Two equivalent implementations: the fast sweep (contiguous-id lattices,
+    i.e. every Lattice built by sprig.dp.lattice) keeps per-level beta blocks
+    in a list and maintains the R-compressed child tables
+        Blo[b,r,k] = logsumexp_B(beta[b,r,B] + logV[k,B])   (Bhi with logW)
+    so the per-level child gathers and writes touch [B,N,R] instead of
+    [B,N,S] — identical values (the compression is exactly the logbmm the
+    reference applies per cut row, hoisted per region) but ~S/R times less
+    index/index_put traffic, which dominates the DP's backward. The reference
+    sweep is kept for foreign/stub lattices."""
+    B, _N, S = term_logits.shape
+    device = ell_leaf.device
+    logV = logV.float()
+    logW = logW.float()
+    U_log = U_logmix.float()
+
+    term_score, log_cont, _ = _leaf_terms(
+        ell_leaf, term_logits, U_log, log_PT, lattice, temper_kappa,
+        term_mark=term_mark, semiring="sum", exact=exact,
+    )
+    tls_flat = _cut_logp_flat(cut_logits, lattice)
+    U_log_t = U_log.transpose(1, 2)
+    rid_mt, slot_mt = _mt_meta(lattice)
+
+    if getattr(lattice, "contiguous_ids", False):
+        return _sweep_fast(
+            term_score, log_cont, cut_logits, U_log_t, logV, logW, lattice,
+            tls_flat, slot_mt, expand_mark, cut_mark, exact,
+        )
+
+    beta = torch.full((B, lattice.n_regions, S), NEG, device=device)
+    beta[:, rid_mt, :] = term_score[:, slot_mt, :]
+
+    row_offset = 0
+    for lv, meta in zip(lattice.levels, _sweep_meta(lattice)):
+        has_leaf, has_nonleaf, leaf_pos, nonleaf_pos, leaf_rids, nonleaf_rids, leaf_slots = meta
+        M = lv.parent_ids.shape[0]
+        P = lv.parents.shape[0]
+        # Bhat[b,m,k] = logsumexp_B(beta[b, child_lo[m], B] + logV[k, B]); Chat same with logW.
+        Bhat = logbmm(beta[:, lv.child_lo, :], logV.t(), exact=exact)   # [B, M, R]
+        Chat = logbmm(beta[:, lv.child_hi, :], logW.t(), exact=exact)
+        contrib = _level_cut_logp(tls_flat, cut_logits, lattice, lv) + Bhat + Chat
+        if cut_mark is not None:
+            contrib = contrib + cut_mark[:, row_offset : row_offset + M].unsqueeze(-1)
+        comb = _segment_logsumexp(contrib, lv.parent_index, P)          # [B, P, R]
+        # expand[b,p,A] = logsumexp_k(comb[b,p,k] + U_log[b,A,k])
+        expand = logbmm_batched(comb, U_log_t, exact=exact)             # [B, P, S]
+        if expand_mark is not None:
+            expand = expand + expand_mark[:, lv.parents, :]
+
+        # Combine the terminate/expand branches. Only finite values may enter
+        # the log-add: with the -1e30 sentinel (and equally with legitimate
+        # log-domain gaps > ~88 nats once emissions sharpen) torch.logaddexp's
+        # first-order backward materializes exp(b - a) = inf nodes whose
+        # double-backward (create_graph=True in SPRIGModel.loss) is NaN — so
+        # mask/index the branches and use a detached-max shift, which keeps
+        # every exp argument <= 0 and the log argument in [1, 2].
+        if has_nonleaf:
+            beta[:, nonleaf_rids, :] = expand[:, nonleaf_pos, :]
+        if has_leaf:
+            a = term_score[:, leaf_slots, :]
+            b = log_cont[:, leaf_slots, :] + expand[:, leaf_pos, :]
+            m = torch.maximum(a, b).detach()
+            beta[:, leaf_rids, :] = m + torch.log(torch.exp(a - m) + torch.exp(b - m))
+        row_offset += M
+
+    logZ = beta[:, lattice.root_id, 0]
+    return beta, logZ
+
+
+def _combine_branches(
+    expand: torch.Tensor,
+    term_score: torch.Tensor,
+    log_cont: torch.Tensor,
+    lv,
+) -> torch.Tensor:
+    """Terminate/expand combination for one level -> beta block [B, P, S] in
+    parent order. Same guarded log-add as the reference sweep (see the long
+    comment there); rows are assembled with a precomputed permutation gather
+    instead of index_put (cheaper backward, no clone of the block)."""
+    if not lv.has_leaf:
+        return expand
+    a = term_score[:, lv.leaf_slots, :]
+    b = log_cont[:, lv.leaf_slots, :] + expand[:, lv.leaf_pos, :]
+    m = torch.maximum(a, b).detach()
+    leaf_val = m + torch.log(torch.exp(a - m) + torch.exp(b - m))
+    if not lv.has_nonleaf:
+        return leaf_val
+    both = torch.cat([expand[:, lv.nonleaf_pos, :], leaf_val], dim=1)
+    return both[:, lv.reorder, :]
+
+
+def _sweep_fast(
+    term_score: torch.Tensor,
+    log_cont: torch.Tensor,
+    cut_logits: torch.Tensor,
+    U_log_t: torch.Tensor,
+    logV: torch.Tensor,
+    logW: torch.Tensor,
+    lattice,
+    tls_flat: Optional[torch.Tensor],
+    slot_mt: torch.Tensor,
+    expand_mark: Optional[torch.Tensor],
+    cut_mark: Optional[torch.Tensor],
+    exact: bool,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Fast inside sweep over contiguous-id lattices (see _sweep docstring)."""
+    B, S = term_score.shape[0], term_score.shape[-1]
+    R = logV.shape[0]
+    N = lattice.n_regions
+    device = term_score.device
+
+    base = term_score[:, slot_mt, :]                       # [B, n_mt, S] == ids [0, n_mt)
+    parts = [base]
+    Blo = torch.full((B, N, R), NEG, device=device)
+    Bhi = torch.full((B, N, R), NEG, device=device)
+    n_mt = base.shape[1]
+    Blo[:, :n_mt, :] = logbmm(base, logV.t(), exact=exact)
+    Bhi[:, :n_mt, :] = logbmm(base, logW.t(), exact=exact)
+
+    row_offset = 0
+    for lv in lattice.levels:
+        M = lv.parent_ids.shape[0]
+        P = lv.parents.shape[0]
+        Bhat = Blo[:, lv.child_lo, :]                      # [B, M, R]
+        Chat = Bhi[:, lv.child_hi, :]
+        contrib = _level_cut_logp(tls_flat, cut_logits, lattice, lv) + Bhat + Chat
+        if cut_mark is not None:
+            contrib = contrib + cut_mark[:, row_offset : row_offset + M].unsqueeze(-1)
+        comb = _segment_logsumexp(contrib, lv.parent_index, P)          # [B, P, R]
+        expand = logbmm_batched(comb, U_log_t, exact=exact)             # [B, P, S]
+        if expand_mark is not None:
+            expand = expand + expand_mark[:, lv.parents, :]
+        val = _combine_branches(expand, term_score, log_cont, lv)       # [B, P, S]
+        parts.append(val)
+        Blo[:, lv.id_lo : lv.id_hi, :] = logbmm(val, logV.t(), exact=exact)
+        Bhi[:, lv.id_lo : lv.id_hi, :] = logbmm(val, logW.t(), exact=exact)
+        row_offset += M
+
+    beta = torch.cat(parts, dim=1)                         # region-id order
+    logZ = beta[:, lattice.root_id, 0]
+    return beta, logZ
+
+
+def inside(
+    ell_leaf: torch.Tensor,
+    term_logits: torch.Tensor,
+    cut_logits: torch.Tensor,
+    U_logmix: torch.Tensor,
+    logV: torch.Tensor,
+    logW: torch.Tensor,
+    lattice,
+    temper_kappa: torch.Tensor,
+    log_PT: torch.Tensor,
+    exact: bool = False,
+) -> Tuple[torch.Tensor, torch.Tensor]:
+    """DESIGN section 5 inside DP -> (beta [B, N_reg, S] fp32, logZ [B])."""
+    return _sweep(
+        ell_leaf, term_logits, cut_logits, U_logmix, logV, logW, lattice,
+        temper_kappa, log_PT, exact=exact,
+    )
+
+
+# DESIGN.md section 5 names the function `inside`; SPRIGModel's adapter looks
+# for `inside_logZ` first — keep both bound to the same implementation.
+inside_logZ = inside
+
+
+def posterior_marginals(
+    ell_leaf: torch.Tensor,
+    term_logits: torch.Tensor,
+    cut_logits: torch.Tensor,
+    U_logmix: torch.Tensor,
+    logV: torch.Tensor,
+    logW: torch.Tensor,
+    lattice,
+    temper_kappa: torch.Tensor,
+    log_PT: torch.Tensor,
+    exact: bool = False,
+) -> Dict[str, torch.Tensor]:
+    """Expected-count posterior marginals via the autograd identity.
+
+    Zero-valued marker potentials are added on the terminate branch, the
+    expand branch, and every concrete cut row; grad of logZ w.r.t. a marker
+    equals the posterior expected count of that event. Returns detached
+    tensors: node/term/expand [B, N_reg, S], cut [B, M_total] (levels
+    concatenated in lattice.levels row order), texel [B, n_leaf, T_v]
+    (grad w.r.t. ell rescaled by kappa), rule [B, S, R] (grad w.r.t.
+    U_logmix), logZ [B].
+    """
+    B, N, S = term_logits.shape
+    device = ell_leaf.device
+    M_total = sum(int(lv.parent_ids.shape[0]) for lv in lattice.levels)
+    ell_d = ell_leaf.detach().requires_grad_(True)
+    U_d = U_logmix.detach().requires_grad_(True)
+    tm = torch.zeros(B, N, S, device=device, requires_grad=True)
+    em = torch.zeros(B, N, S, device=device, requires_grad=True)
+    cm = torch.zeros(B, M_total, device=device, requires_grad=True)
+    with torch.enable_grad():
+        _, logZ = _sweep(
+            ell_d, term_logits.detach(), cut_logits.detach(), U_d,
+            logV.detach(), logW.detach(), lattice,
+            temper_kappa.detach(), log_PT.detach(),
+            term_mark=tm, expand_mark=em, cut_mark=cm, exact=exact,
+        )
+        g_tm, g_em, g_cm, g_ell, g_u = torch.autograd.grad(
+            logZ.sum(), [tm, em, cm, ell_d, U_d]
+        )
+    texel = g_ell * temper_kappa.detach().float().view(1, -1, 1)
+    return {
+        "logZ": logZ.detach(),
+        "node": (g_tm + g_em).detach(),
+        "term": g_tm.detach(),
+        "expand": g_em.detach(),
+        "cut": g_cm.detach(),
+        "texel": texel.detach(),
+        "rule": g_u.detach(),
+    }
+
+
+def viterbi(
+    ell_leaf: torch.Tensor,
+    term_logits: torch.Tensor,
+    cut_logits: torch.Tensor,
+    U_logmix: torch.Tensor,
+    logV: torch.Tensor,
+    logW: torch.Tensor,
+    lattice,
+    temper_kappa: torch.Tensor,
+    log_PT: torch.Tensor,
+    exact: bool = False,
+) -> Tuple[torch.Tensor, List[Tuple]]:
+    """Max-semiring sweep + argmax backtrace.
+
+    Returns (score [B], trees): trees[b] is a nested tuple
+    (region_id, symbol, texel_or_None, axis_or_None, cut_px_or_None,
+    (children...)) rooted at (lattice.root_id, symbol 0). The texel T-mix is
+    a max; the texel prior p(T|A,c) itself stays an exact sum over k.
+    """
+    with torch.no_grad():
+        B, N, S = term_logits.shape
+        R = U_logmix.shape[-1]
+        device = ell_leaf.device
+        U_log = U_logmix.float()
+        logV_f = logV.float()
+        logW_f = logW.float()
+        term_score, log_cont, argT = _leaf_terms(
+            ell_leaf, term_logits, U_log, log_PT, lattice, temper_kappa,
+            semiring="max", exact=exact,
+        )
+        beta = torch.full((B, N, S), NEG, device=device)
+        term_choice = torch.zeros(B, N, S, dtype=torch.bool, device=device)
+        leaf_slot = lattice.leaf_index_of_region
+        rid_mt, slot_mt = _mt_meta(lattice)
+        beta[:, rid_mt, :] = term_score[:, slot_mt, :]
+        term_choice[:, rid_mt, :] = True
+
+        tls_flat = _cut_logp_flat(cut_logits, lattice)
+        metas = _sweep_meta(lattice)
+        bp: List[Dict[str, torch.Tensor]] = []
+        parent_pos: Dict[int, Tuple[int, int]] = {}
+        for li, lv in enumerate(lattice.levels):
+            for p_i, r in enumerate(lv.parents.tolist()):
+                parent_pos[r] = (li, p_i)
+            M = lv.parent_ids.shape[0]
+            P = lv.parents.shape[0]
+
+            # Per-row maxima over child symbols and over components, chunked
+            # over rows to bound the [B, chunk, R, S] / [B, chunk, S, R]
+            # materializations.
+            step = max(1, _MAX_SEMIRING_BUDGET // max(1, B * S * R))
+            cutlp = _level_cut_logp(tls_flat, cut_logits, lattice, lv)   # [B, M, R]
+            argB = torch.empty(B, M, R, dtype=torch.int64, device=device)
+            argC = torch.empty(B, M, R, dtype=torch.int64, device=device)
+            argk = torch.empty(B, M, S, dtype=torch.int64, device=device)
+            tmax = torch.empty(B, M, S, device=device)
+            for m0 in range(0, M, step):
+                m1 = min(M, m0 + step)
+                lo_b = beta[:, lv.child_lo[m0:m1], :]                   # [B, c, S]
+                hi_b = beta[:, lv.child_hi[m0:m1], :]
+                Bv, aB = (lo_b.unsqueeze(2) + logV_f.view(1, 1, R, S)).max(dim=-1)
+                Cv, aC = (hi_b.unsqueeze(2) + logW_f.view(1, 1, R, S)).max(dim=-1)
+                argB[:, m0:m1] = aB
+                argC[:, m0:m1] = aC
+                contrib = cutlp[:, m0:m1] + Bv + Cv                     # [B, c, R]
+                t = contrib.unsqueeze(2) + U_log.unsqueeze(1)           # [B, c, S, R]
+                tm_c, ak = t.max(dim=-1)
+                tmax[:, m0:m1] = tm_c
+                argk[:, m0:m1] = ak
+
+            # Group-max over the rows of each parent; ties -> lowest row index.
+            pidx = lv.parent_index.view(1, -1, 1).expand(B, M, S)
+            pmax = torch.full((B, P, S), NEG, device=device).scatter_reduce(
+                1, pidx, tmax, reduce="amax", include_self=True
+            )
+            rowidx = torch.arange(M, device=device).view(1, -1, 1).expand(B, M, S)
+            cand = torch.where(
+                tmax >= pmax[:, lv.parent_index, :], rowidx,
+                torch.full_like(rowidx, M),
+            )
+            argm = torch.full((B, P, S), M, dtype=torch.int64, device=device).scatter_reduce(
+                1, pidx, cand, reduce="amin", include_self=True
+            )
+
+            rid = lv.parents
+            has_leaf, _hn, leaf_pos, _np, _lr, _nr, leaf_slots = metas[li]
+            cont = torch.zeros(B, P, S, device=device)
+            tb = torch.full((B, P, S), NEG, device=device)
+            if has_leaf:
+                cont[:, leaf_pos, :] = log_cont[:, leaf_slots, :]
+                tb[:, leaf_pos, :] = term_score[:, leaf_slots, :]
+            exp_v = cont + pmax
+            beta[:, rid, :] = torch.maximum(tb, exp_v)
+            term_choice[:, rid, :] = tb >= exp_v
+            bp.append({"argB": argB, "argC": argC, "argk": argk, "argm": argm})
+
+        score = beta[:, lattice.root_id, 0]
+
+        def extract(b: int, rid: int, sym: int) -> Tuple:
+            if bool(term_choice[b, rid, sym]):
+                slot = int(leaf_slot[rid])
+                tex = int(argT[b, slot, sym])
+                return (rid, sym, tex, None, None, ())
+            li, p_i = parent_pos[rid]
+            lv = lattice.levels[li]
+            m = int(bp[li]["argm"][b, p_i, sym])
+            k = int(bp[li]["argk"][b, m, sym])
+            b_sym = int(bp[li]["argB"][b, m, k])
+            c_sym = int(bp[li]["argC"][b, m, k])
+            lo, hi = int(lv.child_lo[m]), int(lv.child_hi[m])
+            axis, px = int(lv.cut_axis[m]), int(lv.cut_px[m])
+            return (
+                rid, sym, None, axis, px,
+                (extract(b, lo, b_sym), extract(b, hi, c_sym)),
+            )
+
+        trees = [extract(b, lattice.root_id, 0) for b in range(B)]
+    return score, trees

sprig/dp/lattice.py

@@ -0,0 +1,286 @@
+"""Finite region lattice for the SPRIG inside DP.
+
+Regions are axis-aligned rectangles whose corners lie on a fixed pixel grid.
+Every cell-interval pair is a region; every interior grid line of a region is
+a valid cut producing two child regions. See DESIGN.md section 3.
+"""
+from __future__ import annotations
+
+import math
+from dataclasses import dataclass, field
+from typing import Dict, List, Optional, Tuple
+
+import torch
+
+N_CUT_TYPES = 14  # 2 axes x 7 relative-offset buckets {1/8 .. 7/8}
+_OFFSET_BUCKETS = [i / 8.0 for i in range(1, 8)]
+
+
+def _bucket(rel: float) -> int:
+    return min(range(7), key=lambda i: abs(_OFFSET_BUCKETS[i] - rel))
+
+
+def cut_type_id(axis: int, rel: float) -> int:
+    """axis: 0 = vertical cut (splits x), 1 = horizontal cut (splits y)."""
+    return axis * 7 + _bucket(rel)
+
+
+@dataclass
+class LatticeLevel:
+    """Flattened cut tables for all parent regions of one cell-area level."""
+    parent_ids: torch.Tensor    # int64 [M]
+    cut_type: torch.Tensor      # int64 [M]  global cut-type id
+    log_cnt: torch.Tensor       # fp32  [M]  log(#same-type cuts in this region)
+    child_lo: torch.Tensor      # int64 [M]  lesser-coordinate child region id
+    child_hi: torch.Tensor      # int64 [M]
+    cut_axis: torch.Tensor      # int64 [M]  0=v, 1=h
+    cut_px: torch.Tensor        # int64 [M]  absolute cut coordinate in px
+    parents: torch.Tensor       # int64 [P]  unique parent ids at this level
+    parent_index: torch.Tensor  # int64 [M]  index into `parents` for scatter
+    # Precomputed sweep metadata (pure functions of the lattice; filled in by
+    # Lattice.__post_init__ so the inside DP never branches on device tensors
+    # or bool-mask-indexes inside the per-level loop — those force CPU-GPU
+    # syncs). Optional so externally-built levels stay constructible;
+    # sprig.dp.inside falls back to computing them on the fly if absent.
+    leaf_pos: Optional[torch.Tensor] = None      # int64 [Pl]  positions in `parents` that are leaf-eligible
+    nonleaf_pos: Optional[torch.Tensor] = None   # int64 [Pn]
+    leaf_rids: Optional[torch.Tensor] = None     # int64 [Pl]  = parents[leaf_pos]
+    nonleaf_rids: Optional[torch.Tensor] = None  # int64 [Pn]
+    leaf_slots: Optional[torch.Tensor] = None    # int64 [Pl]  leaf_index_of_region[leaf_rids]
+    flat_type_idx: Optional[torch.Tensor] = None  # int64 [M] pattern_of_region[parent]*14 + cut_type
+    has_leaf: bool = False
+    has_nonleaf: bool = False
+    reorder: Optional[torch.Tensor] = None       # int64 [P]: cat(nonleaf, leaf) -> parent order
+    id_lo: int = 0                               # parents == arange(id_lo, id_hi) when the
+    id_hi: int = 0                               # lattice has contiguous_ids (see Lattice)
+
+
+@dataclass
+class Lattice:
+    canvas_px: int
+    grid_px: int
+    leaf_max_px: int
+
+    regions: torch.Tensor = field(init=False)         # int64 [N,4] x0,y0,x1,y1 (px)
+    region_id: Dict[Tuple[int, int, int, int], int] = field(init=False)
+    leaf_mask: torch.Tensor = field(init=False)        # bool [N]
+    must_terminate: torch.Tensor = field(init=False)   # bool [N]
+    must_expand: torch.Tensor = field(init=False)      # bool [N]
+    area_px: torch.Tensor = field(init=False)          # fp32 [N]
+    phi_geom: torch.Tensor = field(init=False)         # fp32 [N,64]
+    type_present: torch.Tensor = field(init=False)     # bool [N,14]
+    levels: List[LatticeLevel] = field(init=False)     # area-ascending, only levels with cuts
+    root_id: int = field(init=False)
+    leaf_ids: torch.Tensor = field(init=False)         # int64 [NL] leaf-eligible region ids
+    leaf_index_of_region: torch.Tensor = field(init=False)  # int64 [N] (-1 if not leaf)
+    cuts_of_region: Dict[int, List[Tuple[int, int, int, int, int, float]]] = field(init=False)
+    # cuts_of_region[r] = list of (axis, px, child_lo, child_hi, type_id, log_cnt)
+    # Precomputed sweep/emission metadata (see __post_init__).
+    mt_ids: torch.Tensor = field(init=False)           # int64 [n_mt] must-terminate region ids
+    mt_slots: torch.Tensor = field(init=False)         # int64 [n_mt] their leaf slots
+    type_patterns: torch.Tensor = field(init=False)    # bool [n_pat, 14] distinct type_present rows
+    pattern_of_region: torch.Tensor = field(init=False)  # int64 [N] (-1 if region has no cuts)
+    shape_groups: List[Tuple[int, int, torch.Tensor, torch.Tensor]] = field(init=False)
+    contiguous_ids: bool = field(init=False)           # id blocks are (mt, level0, level1, ...)
+
+    def __post_init__(self) -> None:
+        g, c = self.grid_px, self.canvas_px
+        assert c % g == 0
+        n_cells = c // g
+        lines = [i * g for i in range(n_cells + 1)]
+
+        rects: List[Tuple[int, int, int, int]] = []
+        for x0 in range(n_cells):
+            for x1 in range(x0 + 1, n_cells + 1):
+                for y0 in range(n_cells):
+                    for y1 in range(y0 + 1, n_cells + 1):
+                        rects.append((lines[x0], lines[y0], lines[x1], lines[y1]))
+        rects.sort(key=lambda r: ((r[2] - r[0]) * (r[3] - r[1]), r))
+        self.region_id = {r: i for i, r in enumerate(rects)}
+        self.regions = torch.tensor(rects, dtype=torch.int64)
+        N = len(rects)
+
+        w = self.regions[:, 2] - self.regions[:, 0]
+        h = self.regions[:, 3] - self.regions[:, 1]
+        self.leaf_mask = (w <= self.leaf_max_px) & (h <= self.leaf_max_px)
+        self.must_terminate = (w == g) & (h == g)
+        self.must_expand = (w > self.leaf_max_px) | (h > self.leaf_max_px)
+        self.area_px = (w * h).to(torch.float32)
+        self.root_id = self.region_id[(0, 0, c, c)]
+
+        self.leaf_ids = torch.nonzero(self.leaf_mask, as_tuple=False).squeeze(1)
+        self.leaf_index_of_region = torch.full((N,), -1, dtype=torch.int64)
+        self.leaf_index_of_region[self.leaf_ids] = torch.arange(len(self.leaf_ids))
+
+        # Fourier geometry features (log-area, log-aspect, center x/y) -> 64 dims.
+        la = torch.log(self.area_px / (c * c))
+        lasp = torch.log(w.float() / h.float())
+        cx = (self.regions[:, 0] + self.regions[:, 2]).float() / (2 * c)
+        cy = (self.regions[:, 1] + self.regions[:, 3]).float() / (2 * c)
+        base = torch.stack([la / 8.0, lasp / 4.0, cx, cy], dim=1)  # [N,4]
+        freqs = torch.tensor([1.0, 2.0, 4.0, 8.0, 16.0, 32.0, 3.0, 5.0])  # 8 freqs
+        ang = base.unsqueeze(2) * freqs.view(1, 1, 8) * math.pi  # [N,4,8]
+        self.phi_geom = torch.cat([torch.sin(ang), torch.cos(ang)], dim=2).reshape(N, 64)
+
+        # Enumerate cuts.
+        self.cuts_of_region = {}
+        self.type_present = torch.zeros(N, N_CUT_TYPES, dtype=torch.bool)
+        per_level: Dict[int, List[Tuple[int, int, int, float, int, int, int]]] = {}
+        for rid, (x0, y0, x1, y1) in enumerate(rects):
+            cuts: List[Tuple[int, int, int, int, int, float]] = []
+            type_counts: Dict[int, int] = {}
+            cand: List[Tuple[int, int, int, int, int]] = []  # axis, px, lo, hi, type
+            for px in range(x0 + g, x1, g):        # vertical cuts
+                lo = self.region_id[(x0, y0, px, y1)]
+                hi = self.region_id[(px, y0, x1, y1)]
+                t = cut_type_id(0, (px - x0) / (x1 - x0))
+                cand.append((0, px, lo, hi, t))
+                type_counts[t] = type_counts.get(t, 0) + 1
+            for py in range(y0 + g, y1, g):        # horizontal cuts
+                lo = self.region_id[(x0, y0, x1, py)]
+                hi = self.region_id[(x0, py, x1, y1)]
+                t = cut_type_id(1, (py - y0) / (y1 - y0))
+                cand.append((1, py, lo, hi, t))
+                type_counts[t] = type_counts.get(t, 0) + 1
+            for axis, px, lo, hi, t in cand:
+                lc = math.log(type_counts[t])
+                cuts.append((axis, px, lo, hi, t, lc))
+                self.type_present[rid, t] = True
+            self.cuts_of_region[rid] = cuts
+            if cuts and not bool(self.must_terminate[rid]):
+                area_cells = int(self.area_px[rid].item()) // (g * g)
+                per_level.setdefault(area_cells, [])
+                for axis, px, lo, hi, t, lc in cuts:
+                    per_level[area_cells].append((rid, t, lo, hi, lc, axis, px))
+
+        self.levels = []
+        for area_cells in sorted(per_level.keys()):
+            rows = per_level[area_cells]
+            parent_ids = torch.tensor([r[0] for r in rows], dtype=torch.int64)
+            parents, parent_index = torch.unique(parent_ids, return_inverse=True)
+            self.levels.append(LatticeLevel(
+                parent_ids=parent_ids,
+                cut_type=torch.tensor([r[1] for r in rows], dtype=torch.int64),
+                log_cnt=torch.tensor([r[4] for r in rows], dtype=torch.float32),
+                child_lo=torch.tensor([r[2] for r in rows], dtype=torch.int64),
+                child_hi=torch.tensor([r[3] for r in rows], dtype=torch.int64),
+                cut_axis=torch.tensor([r[5] for r in rows], dtype=torch.int64),
+                cut_px=torch.tensor([r[6] for r in rows], dtype=torch.int64),
+                parents=parents,
+                parent_index=parent_index,
+            ))
+
+        # ---- precomputed sweep metadata (everything shape-related the DP
+        # would otherwise derive from device tensors inside the level loop).
+        self.mt_ids = torch.nonzero(self.must_terminate, as_tuple=False).reshape(-1)
+        self.mt_slots = self.leaf_index_of_region[self.mt_ids]
+
+        # Distinct cut-type-presence patterns over regions that actually have
+        # cuts (every level parent). Lets the DP do ONE masked log-softmax
+        # over [B, n_pat, R, 14] per sweep instead of one per level.
+        pat_of: Dict[Tuple[bool, ...], int] = {}
+        pats: List[torch.Tensor] = []
+        self.pattern_of_region = torch.full((N,), -1, dtype=torch.int64)
+        for lv in self.levels:
+            for rid in lv.parents.tolist():
+                key = tuple(self.type_present[rid].tolist())
+                pid = pat_of.setdefault(key, len(pats))
+                if pid == len(pats):
+                    pats.append(self.type_present[rid].clone())
+                self.pattern_of_region[rid] = pid
+        self.type_patterns = torch.stack(pats, dim=0)          # bool [n_pat, 14]
+
+        for lv in self.levels:
+            is_leaf = self.leaf_mask[lv.parents]
+            lv.leaf_pos = torch.nonzero(is_leaf, as_tuple=False).reshape(-1)
+            lv.nonleaf_pos = torch.nonzero(~is_leaf, as_tuple=False).reshape(-1)
+            lv.leaf_rids = lv.parents[lv.leaf_pos]
+            lv.nonleaf_rids = lv.parents[lv.nonleaf_pos]
+            lv.leaf_slots = self.leaf_index_of_region[lv.leaf_rids]
+            lv.flat_type_idx = (
+                self.pattern_of_region[lv.parent_ids] * N_CUT_TYPES + lv.cut_type
+            )
+            lv.has_leaf = bool(lv.leaf_pos.numel() > 0)
+            lv.has_nonleaf = bool(lv.nonleaf_pos.numel() > 0)
+            lv.reorder = torch.argsort(torch.cat([lv.nonleaf_pos, lv.leaf_pos]))
+
+        # Region ids are assigned area-ascending (rects sorted by pixel area),
+        # so the must-terminate block and each level's parent block are
+        # contiguous, consecutive id ranges covering all N regions. The DP's
+        # fast sweep relies on this to assemble beta with one cat instead of
+        # per-level index_put on the [B, N, S] table. Verified here; any
+        # violation falls back to the reference sweep.
+        self.contiguous_ids = bool(
+            self.mt_ids.equal(torch.arange(self.mt_ids.numel()))
+        )
+        off = int(self.mt_ids.numel())
+        for lv in self.levels:
+            P = int(lv.parents.numel())
+            if not bool(lv.parents.equal(torch.arange(off, off + P))):
+                self.contiguous_ids = False
+                break
+            lv.id_lo, lv.id_hi = off, off + P
+            off += P
+        if off != N:
+            self.contiguous_ids = False
+
+        # Leaf shape groups with flattened pixel-gather indices, cached for
+        # emission scoring (atlas.score_leaves): (h, w, slots [n_g],
+        # pix [n_g*h*w] into a [canvas*canvas]-flattened image).
+        self.shape_groups = []
+        for (h, w), slots in self.leaf_shape_groups().items():
+            rects = self.regions[self.leaf_ids[slots]]
+            yy = rects[:, 1].view(-1, 1, 1) + torch.arange(h).view(1, -1, 1)
+            xx = rects[:, 0].view(-1, 1, 1) + torch.arange(w).view(1, 1, -1)
+            pix = (yy * c + xx).reshape(-1)
+            self.shape_groups.append((h, w, slots, pix))
+
+    @property
+    def n_regions(self) -> int:
+        return int(self.regions.shape[0])
+
+    @property
+    def n_leaf_regions(self) -> int:
+        return int(self.leaf_ids.shape[0])
+
+    def leaf_shape_groups(self) -> Dict[Tuple[int, int], torch.Tensor]:
+        """Map (h, w) -> leaf-slot indices (into the leaf_ids ordering)."""
+        groups: Dict[Tuple[int, int], List[int]] = {}
+        for slot, rid in enumerate(self.leaf_ids.tolist()):
+            x0, y0, x1, y1 = self.regions[rid].tolist()
+            groups.setdefault((y1 - y0, x1 - x0), []).append(slot)
+        return {k: torch.tensor(v, dtype=torch.int64) for k, v in groups.items()}
+
+    def to(self, device: torch.device) -> "Lattice":
+        device = torch.device(device)
+        if getattr(self, "_device", None) == device:
+            return self
+        for name in ("regions", "leaf_mask", "must_terminate", "must_expand",
+                     "area_px", "phi_geom", "type_present", "leaf_ids",
+                     "leaf_index_of_region", "mt_ids", "mt_slots",
+                     "type_patterns", "pattern_of_region"):
+            setattr(self, name, getattr(self, name).to(device))
+        for lv in self.levels:
+            for name in ("parent_ids", "cut_type", "log_cnt", "child_lo", "child_hi",
+                         "cut_axis", "cut_px", "parents", "parent_index",
+                         "leaf_pos", "nonleaf_pos", "leaf_rids", "nonleaf_rids",
+                         "leaf_slots", "flat_type_idx", "reorder"):
+                t = getattr(lv, name)
+                if t is not None:
+                    setattr(lv, name, t.to(device))
+        self.shape_groups = [
+            (h, w, slots.to(device), pix.to(device))
+            for h, w, slots, pix in self.shape_groups
+        ]
+        self._device = device
+        return self
+
+
+_CACHE: Dict[Tuple[int, int, int], Lattice] = {}
+
+
+def get_lattice(canvas_px: int, grid_px: int, leaf_max_px: int) -> Lattice:
+    key = (canvas_px, grid_px, leaf_max_px)
+    if key not in _CACHE:
+        _CACHE[key] = Lattice(canvas_px, grid_px, leaf_max_px)
+    return _CACHE[key]

sprig/eval/baseline_pixmix.py

@@ -0,0 +1,88 @@
+"""B0 trivial baseline: position-dependent 4-component DL mixture per pixel.
+
+One parameter grid [64, 64, 40] (no neural net, no caption): per pixel, the
+same 40-channel discretized-logistic mixture parameterization as the texel
+atlas (sprig/model/dl.py layout: 4 components x contiguous blocks of
+[weight, 3 means, 3 log-scales, 3 coupling coeffs]). B0's bpd measures
+exactly what the grammar buys over a caption-free per-pixel model.
+"""
+from __future__ import annotations
+
+import math
+from typing import Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+
+from sprig.model.dl import LOGSCALE_IDX, N_CH, N_COMP, dl_logprob, u8_to_unit
+
+LOG2 = math.log(2.0)
+N_PARAMS = N_CH  # 40
+
+
+def _to_u8_tensor(images_u8: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
+    if isinstance(images_u8, np.ndarray):
+        images_u8 = torch.from_numpy(np.array(images_u8, dtype=images_u8.dtype))
+    return images_u8
+
+
+class PixMixBaseline(nn.Module):
+    """params: [H, W, 40] trainable grid; no other parameters."""
+
+    def __init__(self, height: int = 64, width: int = 64):
+        super().__init__()
+        p = torch.zeros(height, width, N_PARAMS)
+        # spread component means so the mixture starts multi-modal
+        comp_means = torch.linspace(-0.6, 0.6, N_COMP)
+        for j in range(N_COMP):
+            p[..., 10 * j + 1: 10 * j + 4] = comp_means[j]
+        p[..., LOGSCALE_IDX] = -1.0
+        self.params = nn.Parameter(p)
+
+    def logprob_per_image(self, images_u8: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
+        """[B] total log-likelihood (nats) per image."""
+        img = _to_u8_tensor(images_u8).to(self.params.device)
+        x = u8_to_unit(img).permute(0, 3, 1, 2)               # [B,3,H,W]
+        # expand (view) params to the batch lead dim: dl_logprob broadcasts
+        # x against params but not vice versa
+        p = self.params.permute(2, 0, 1).unsqueeze(0).expand(x.shape[0], -1, -1, -1)
+        lp = dl_logprob(p, x)                                 # [B,H,W] fp32
+        return lp.sum(dim=(1, 2))
+
+    @torch.no_grad()
+    def bpd(self, images_u8: Union[np.ndarray, torch.Tensor], batch_size: int = 256) -> float:
+        """Mean bits/dim over images (dims = 3*H*W)."""
+        n = int(images_u8.shape[0])
+        h, w = self.params.shape[0], self.params.shape[1]
+        total = 0.0
+        for i in range(0, n, batch_size):
+            total += float(self.logprob_per_image(images_u8[i:i + batch_size]).sum())
+        return -total / (n * 3.0 * h * w * LOG2)
+
+
+def fit(
+    images_u8: Union[np.ndarray, torch.Tensor],
+    steps: int = 2000,
+    lr: float = 5e-2,
+    batch_size: int = 128,
+    device: str = "cpu",
+    seed: int = 0,
+) -> PixMixBaseline:
+    """Fit B0 by Adam on random minibatches of training images; returns the model."""
+    torch.manual_seed(seed)
+    model = PixMixBaseline().to(device)
+    opt = torch.optim.Adam(model.parameters(), lr=lr)
+    n = images_u8.shape[0]
+    gen = np.random.default_rng(seed)
+    for _ in range(steps):
+        idx = gen.integers(0, n, size=min(batch_size, n))
+        if isinstance(images_u8, np.ndarray):
+            batch = images_u8[idx]
+        else:
+            batch = images_u8[torch.from_numpy(idx)]
+        nll = -model.logprob_per_image(batch).mean() / (3.0 * 64 * 64)
+        opt.zero_grad()
+        nll.backward()
+        opt.step()
+    return model

sprig/eval/color_checks.py

@@ -0,0 +1,189 @@
+"""GT-free faithfulness checks on generated images.
+
+Pipeline (pure numpy, no scipy):
+background = modal border color; foreground = pixels whose max-channel
+deviation from the background exceeds `thresh`; objects = 4-connected
+components of the foreground; object color = median RGB mapped to the
+nearest vocabulary color anchor in CIE Lab; relations compared by centroid.
+"""
+from __future__ import annotations
+
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+
+# Vocabulary color anchors (RGB). Single source of truth is
+# sprig/data/procgen/vocab.py (COLORS dict); these literals are the fallback.
+_DEFAULT_ANCHORS: Dict[str, Tuple[int, int, int]] = {
+    "red": (210, 45, 45),
+    "green": (55, 170, 60),
+    "blue": (50, 90, 215),
+    "yellow": (235, 215, 55),
+    "orange": (240, 145, 40),
+    "purple": (140, 60, 185),
+    "cyan": (60, 205, 215),
+    "magenta": (220, 70, 175),
+}
+
+
+def _load_anchors() -> Dict[str, Tuple[int, int, int]]:
+    try:
+        from sprig.data.procgen import vocab as _vocab  # type: ignore
+
+        for name in ("COLOR_ANCHORS", "COLORS", "COLOR_RGB"):
+            anchors = getattr(_vocab, name, None)
+            if isinstance(anchors, dict) and anchors:
+                first = next(iter(anchors.values()))
+                if hasattr(first, "__len__") and len(first) == 3:
+                    return {str(k): tuple(int(v) for v in rgb) for k, rgb in anchors.items()}
+    except Exception:
+        pass
+    return dict(_DEFAULT_ANCHORS)
+
+
+COLOR_ANCHORS: Dict[str, Tuple[int, int, int]] = _load_anchors()
+
+
+def rgb_to_lab(rgb: np.ndarray) -> np.ndarray:
+    """sRGB (0..255, any leading shape [..,3]) -> CIE Lab (D65), numpy only."""
+    c = np.asarray(rgb, dtype=np.float64) / 255.0
+    lin = np.where(c <= 0.04045, c / 12.92, ((c + 0.055) / 1.055) ** 2.4)
+    m = np.array(
+        [
+            [0.4124564, 0.3575761, 0.1804375],
+            [0.2126729, 0.7151522, 0.0721750],
+            [0.0193339, 0.1191920, 0.9503041],
+        ]
+    )
+    xyz = lin @ m.T
+    xyz = xyz / np.array([0.95047, 1.0, 1.08883])
+    d = 6.0 / 29.0
+    f = np.where(xyz > d ** 3, np.cbrt(xyz), xyz / (3 * d * d) + 4.0 / 29.0)
+    lab = np.empty_like(f)
+    lab[..., 0] = 116.0 * f[..., 1] - 16.0
+    lab[..., 1] = 500.0 * (f[..., 0] - f[..., 1])
+    lab[..., 2] = 200.0 * (f[..., 1] - f[..., 2])
+    return lab
+
+
+_ANCHOR_NAMES = list(COLOR_ANCHORS.keys())
+_ANCHOR_LAB = rgb_to_lab(np.array([COLOR_ANCHORS[n] for n in _ANCHOR_NAMES], dtype=np.float64))
+
+
+def nearest_color(rgb) -> str:
+    """Nearest vocabulary anchor to an RGB triple, in Lab space."""
+    lab = rgb_to_lab(np.asarray(rgb, dtype=np.float64))
+    d = np.linalg.norm(_ANCHOR_LAB - lab[None, :], axis=1)
+    return _ANCHOR_NAMES[int(np.argmin(d))]
+
+
+def _modal_border_color(img: np.ndarray) -> np.ndarray:
+    border = np.concatenate([img[0], img[-1], img[1:-1, 0], img[1:-1, -1]], axis=0)
+    codes = (
+        border[:, 0].astype(np.int64) * 65536
+        + border[:, 1].astype(np.int64) * 256
+        + border[:, 2].astype(np.int64)
+    )
+    vals, counts = np.unique(codes, return_counts=True)
+    code = int(vals[np.argmax(counts)])
+    return np.array([code // 65536, (code // 256) % 256, code % 256], dtype=np.float64)
+
+
+def connected_components(mask: np.ndarray) -> Tuple[np.ndarray, int]:
+    """4-connectivity components of a boolean mask via label propagation.
+
+    Returns (labels [H,W] int32 with 0 = background, 1..n components, n).
+    """
+    h, w = mask.shape
+    labels = np.where(mask, np.arange(1, h * w + 1, dtype=np.int64).reshape(h, w), 0)
+    while True:
+        prev = labels
+        up = np.zeros_like(labels)
+        up[1:, :] = labels[:-1, :]
+        down = np.zeros_like(labels)
+        down[:-1, :] = labels[1:, :]
+        left = np.zeros_like(labels)
+        left[:, 1:] = labels[:, :-1]
+        right = np.zeros_like(labels)
+        right[:, :-1] = labels[:, 1:]
+        stacked = np.stack([labels, up, down, left, right], axis=0)
+        stacked = np.where(stacked == 0, np.iinfo(np.int64).max, stacked)
+        labels = np.where(mask, stacked.min(axis=0), 0)
+        if np.array_equal(labels, prev):
+            break
+    vals = np.unique(labels)
+    vals = vals[vals > 0]
+    out = np.zeros_like(labels, dtype=np.int32)
+    for i, v in enumerate(vals):
+        out[labels == v] = i + 1
+    return out, len(vals)
+
+
+def extract(image: np.ndarray, thresh: int = 40, min_area: int = 12) -> Dict[str, object]:
+    """Extract background + object list from a rendered/generated 64x64 image.
+
+    Returns {"background": color_name, "background_rgb": (r,g,b),
+             "objects": [{"color", "rgb", "centroid": (x,y), "area",
+                          "bbox": (x0,y0,x1,y1)}]} sorted by area descending.
+    """
+    img = np.asarray(image)
+    if img.dtype != np.uint8:
+        img = np.clip(img, 0, 255).astype(np.uint8)
+    bg = _modal_border_color(img)
+    dev = np.abs(img.astype(np.float64) - bg[None, None, :]).max(axis=-1)
+    mask = dev > float(thresh)
+    labels, n = connected_components(mask)
+    objects: List[Dict[str, object]] = []
+    for i in range(1, n + 1):
+        ys, xs = np.nonzero(labels == i)
+        if ys.size < min_area:
+            continue
+        med = np.median(img[ys, xs].astype(np.float64), axis=0)
+        objects.append(
+            {
+                "color": nearest_color(med),
+                "rgb": tuple(float(v) for v in med),
+                "centroid": (float(xs.mean()), float(ys.mean())),
+                "area": int(ys.size),
+                "bbox": (int(xs.min()), int(ys.min()), int(xs.max()) + 1, int(ys.max()) + 1),
+            }
+        )
+    objects.sort(key=lambda o: -int(o["area"]))
+    return {
+        "background": nearest_color(bg),
+        "background_rgb": tuple(float(v) for v in bg),
+        "objects": objects,
+    }
+
+
+def _find_by_color(objects: List[Dict[str, object]], color: str) -> Optional[Dict[str, object]]:
+    for o in objects:
+        if o["color"] == color:
+            return o
+    return None
+
+
+def relation_holds(
+    extraction: Dict[str, object], color_a: str, color_b: str, relation: str
+) -> Optional[bool]:
+    """Does 'the color_a object is <relation> the color_b object' hold by centroids?
+
+    relation in {"left of", "right of", "above", "below"}.
+    Returns None when either object is missing (unscoreable).
+    """
+    objs = extraction["objects"]  # type: ignore[index]
+    a = _find_by_color(objs, color_a)  # type: ignore[arg-type]
+    b = _find_by_color(objs, color_b)  # type: ignore[arg-type]
+    if a is None or b is None:
+        return None
+    ax, ay = a["centroid"]  # type: ignore[misc]
+    bx, by = b["centroid"]  # type: ignore[misc]
+    if relation == "left of":
+        return bool(ax < bx)
+    if relation == "right of":
+        return bool(ax > bx)
+    if relation == "above":
+        return bool(ay < by)
+    if relation == "below":
+        return bool(ay > by)
+    raise ValueError("unknown relation: {}".format(relation))

sprig/eval/monitors.py

@@ -0,0 +1,102 @@
+"""Training-health monitors: pure functions over posterior-usage statistics.
+
+Inputs come from `SPRIGModel.posterior_usage` (DESIGN §4/C3):
+symbol_usage [S], texel_usage [T_v] (expected-count vectors, any positive
+scale — normalized here), node_entropy (nats/node), emit_mag / rule_mag
+(gradient-magnitude scalars), mean_depth, mean_leaves.
+"""
+from __future__ import annotations
+
+from typing import Dict, Tuple, Union
+
+import numpy as np
+
+ArrayLike = Union[np.ndarray, "torch.Tensor", list]  # noqa: F821
+
+# Alarm thresholds (plan §Verification).
+S_EFF_COLLAPSE_FRAC = 0.15        # S_eff < 0.15*S -> grammar collapse
+MAG_RATIO_ALARM = 100.0           # emission/rule ratio > 100 ...
+NODE_ENTROPY_FLOOR = 0.1          # ... with node entropy < 0.1 nats
+ALIVE_TEXEL_ALARM_FRAC = 0.25     # alive fraction below this -> texel death
+QUADTREE_DEPTH = 6.0              # uniform 64px -> 8px binary-split depth
+QUADTREE_LEAVES = 64.0            # 8x8 grid of 8px leaves
+ALIVE_USAGE_THRESH_FRAC = 0.1     # texel alive iff usage > 0.1 / T_v
+
+
+def _to_numpy(x: ArrayLike) -> np.ndarray:
+    if hasattr(x, "detach"):  # torch tensor without importing torch
+        x = x.detach().cpu().numpy()
+    return np.asarray(x, dtype=np.float64)
+
+
+def s_eff(symbol_usage: ArrayLike) -> float:
+    """Effective symbol count exp(H(usage)); usage is normalized internally."""
+    p = _to_numpy(symbol_usage)
+    total = p.sum()
+    if total <= 0:
+        return 0.0
+    p = p / total
+    nz = p[p > 0]
+    return float(np.exp(-(nz * np.log(nz)).sum()))
+
+
+def alive_texels(texel_usage: ArrayLike, T_v: int) -> int:
+    """Number of texels with normalized usage above 0.1/T_v (resurrection thresh)."""
+    p = _to_numpy(texel_usage)
+    total = p.sum()
+    if total <= 0:
+        return 0
+    p = p / total
+    return int((p > ALIVE_USAGE_THRESH_FRAC / float(T_v)).sum())
+
+
+def magnitude_ratio(emit_mag: float, rule_mag: float, eps: float = 1e-12) -> float:
+    """Emission-vs-rule gradient magnitude ratio (tempered-DP failure detector)."""
+    return float(emit_mag) / (float(rule_mag) + eps)
+
+
+def pi_controller_eta(
+    eta: float,
+    node_entropy: float,
+    band: Tuple[float, float] = (0.5, 3.0),
+    eta_range: Tuple[float, float] = (0.0, 1.5),
+) -> float:
+    """One PI-controller update of the tempering exponent eta (DESIGN §5).
+
+    If posterior node entropy H is below the band, raise eta by
+    0.05 + 0.1*(lo - H); if above the band, lower it by 0.05; clamp to range.
+    """
+    lo, hi = band
+    h = float(node_entropy)
+    new_eta = float(eta)
+    if h < lo:
+        new_eta += 0.05 + 0.1 * (lo - h)
+    elif h > hi:
+        new_eta -= 0.05
+    return float(min(max(new_eta, eta_range[0]), eta_range[1]))
+
+
+def build_alarms(stats: Dict[str, object], S: int, T_v: int) -> Dict[str, bool]:
+    """Alarm dict per the plan's training alarms.
+
+    `stats` is the posterior_usage dict: keys symbol_usage, texel_usage,
+    node_entropy, emit_mag, rule_mag, mean_depth, mean_leaves.
+    Returns booleans: grammar_collapse, tempered_dp_failure, texel_death,
+    parse_collapse, and `any`.
+    """
+    seff = s_eff(stats["symbol_usage"])
+    ratio = magnitude_ratio(float(stats["emit_mag"]), float(stats["rule_mag"]))
+    node_h = float(stats["node_entropy"])
+    alive_frac = alive_texels(stats["texel_usage"], T_v) / float(T_v)
+    mean_depth = float(stats.get("mean_depth", 0.0))
+    mean_leaves = float(stats.get("mean_leaves", 0.0))
+
+    alarms = {
+        "grammar_collapse": seff < S_EFF_COLLAPSE_FRAC * S,
+        "tempered_dp_failure": (ratio > MAG_RATIO_ALARM) and (node_h < NODE_ENTROPY_FLOOR),
+        "texel_death": alive_frac < ALIVE_TEXEL_ALARM_FRAC,
+        "parse_collapse": (mean_leaves >= 0.95 * QUADTREE_LEAVES)
+        and (mean_depth >= 0.95 * QUADTREE_DEPTH),
+    }
+    alarms["any"] = bool(any(alarms.values()))
+    return alarms

sprig/eval/probe.py

@@ -0,0 +1,150 @@
+"""3-conv-block CNN probe: shape (8-way) + color (8-way) heads.
+
+Used for the compositional-holdout score: trained on a probe dataset that
+DOES include the held-out combos (leakage-safe by directory separation),
+then applied to model generations.
+
+Labeled image dir formats accepted by `train` (checked in this order):
+1. memmap pair: `images.u8` [N,64,64,3] + `labels.npy` int64 [N,2]
+   (shape_idx, color_idx) in the canonical prompts.SHAPES / prompts.COLORS order;
+2. flat PNG files named `<shape>_<color>_<anything>.png`.
+"""
+from __future__ import annotations
+
+import glob
+import os
+from typing import List, Optional, Tuple, Union
+
+import numpy as np
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sprig.eval.prompts import COLORS, SHAPES
+
+
+class ProbeCNN(nn.Module):
+    def __init__(self, width: int = 32):
+        super().__init__()
+        w = width
+        self.blocks = nn.Sequential(
+            nn.Conv2d(3, w, 3, padding=1), nn.BatchNorm2d(w), nn.ReLU(),
+            nn.Conv2d(w, w, 3, padding=1), nn.BatchNorm2d(w), nn.ReLU(),
+            nn.MaxPool2d(2),                                                     # 32
+            nn.Conv2d(w, 2 * w, 3, padding=1), nn.BatchNorm2d(2 * w), nn.ReLU(),
+            nn.Conv2d(2 * w, 2 * w, 3, padding=1), nn.BatchNorm2d(2 * w), nn.ReLU(),
+            nn.MaxPool2d(2),                                                     # 16
+            nn.Conv2d(2 * w, 4 * w, 3, padding=1), nn.BatchNorm2d(4 * w), nn.ReLU(),
+            nn.Conv2d(4 * w, 4 * w, 3, padding=1), nn.BatchNorm2d(4 * w), nn.ReLU(),
+            nn.MaxPool2d(2),                                                     # 8
+        )
+        self.shape_head = nn.Linear(4 * w, len(SHAPES))
+        self.color_head = nn.Linear(4 * w, len(COLORS))
+
+    def forward(self, x: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+        h = self.blocks(x)
+        h = h.mean(dim=(2, 3))
+        return self.shape_head(h), self.color_head(h)
+
+
+def _images_to_tensor(images: Union[np.ndarray, torch.Tensor]) -> torch.Tensor:
+    """u8 [N,64,64,3] (numpy or torch) -> float [N,3,64,64] in [-1,1]."""
+    if isinstance(images, np.ndarray):
+        images = torch.from_numpy(np.array(images, dtype=images.dtype))
+    x = images.to(torch.float32) / 127.5 - 1.0
+    return x.permute(0, 3, 1, 2).contiguous()
+
+
+def _load_probe_dir(data_dir: str) -> Tuple[np.ndarray, np.ndarray]:
+    img_path = os.path.join(data_dir, "images.u8")
+    lab_path = os.path.join(data_dir, "labels.npy")
+    if os.path.exists(img_path) and os.path.exists(lab_path):
+        labels = np.load(lab_path).astype(np.int64)
+        n = labels.shape[0]
+        images = np.memmap(img_path, dtype=np.uint8, mode="r", shape=(n, 64, 64, 3))
+        return np.asarray(images), labels
+    files = sorted(glob.glob(os.path.join(data_dir, "*.png")))
+    if not files:
+        raise FileNotFoundError(
+            "probe dir {} has neither images.u8+labels.npy nor *.png".format(data_dir)
+        )
+    from PIL import Image
+
+    imgs: List[np.ndarray] = []
+    labs: List[Tuple[int, int]] = []
+    for f in files:
+        parts = os.path.basename(f).split("_")
+        shape_name, color_name = parts[0], parts[1]
+        labs.append((SHAPES.index(shape_name), COLORS.index(color_name)))
+        imgs.append(np.asarray(Image.open(f).convert("RGB"), dtype=np.uint8))
+    return np.stack(imgs), np.asarray(labs, dtype=np.int64)
+
+
+def train(
+    data_dir: str,
+    epochs: int = 10,
+    batch_size: int = 64,
+    lr: float = 1e-3,
+    device: str = "cpu",
+    ckpt_path: Optional[str] = None,
+    seed: int = 0,
+) -> str:
+    """Train the probe on a labeled image dir; returns the checkpoint path."""
+    images, labels = _load_probe_dir(data_dir)
+    torch.manual_seed(seed)
+    model = ProbeCNN().to(device)
+    opt = torch.optim.Adam(model.parameters(), lr=lr)
+    n = images.shape[0]
+    gen = np.random.default_rng(seed)
+    model.train()
+    for _ in range(epochs):
+        order = gen.permutation(n)
+        for i in range(0, n, batch_size):
+            idx = order[i:i + batch_size]
+            x = _images_to_tensor(images[idx]).to(device)
+            ys = torch.from_numpy(labels[idx, 0]).to(device)
+            yc = torch.from_numpy(labels[idx, 1]).to(device)
+            logit_s, logit_c = model(x)
+            loss = F.cross_entropy(logit_s, ys) + F.cross_entropy(logit_c, yc)
+            opt.zero_grad()
+            loss.backward()
+            opt.step()
+    if ckpt_path is None:
+        ckpt_path = os.path.join(data_dir, "probe.pt")
+    torch.save(
+        {"state_dict": model.state_dict(), "shapes": SHAPES, "colors": COLORS}, ckpt_path
+    )
+    return ckpt_path
+
+
+def load_probe(ckpt_path: str, device: str = "cpu") -> ProbeCNN:
+    ckpt = torch.load(ckpt_path, map_location=device)
+    model = ProbeCNN().to(device)
+    model.load_state_dict(ckpt["state_dict"])
+    model.eval()
+    return model
+
+
+@torch.no_grad()
+def score_generations(
+    images: Union[np.ndarray, torch.Tensor],
+    expected_shape: str,
+    expected_color: str,
+    ckpt_path: Union[str, ProbeCNN],
+    device: str = "cpu",
+    batch_size: int = 64,
+) -> float:
+    """Fraction of images classified as (expected_shape AND expected_color)."""
+    model = ckpt_path if isinstance(ckpt_path, ProbeCNN) else load_probe(ckpt_path, device)
+    model.eval()
+    target_s = SHAPES.index(expected_shape)
+    target_c = COLORS.index(expected_color)
+    x_all = _images_to_tensor(images)
+    hits = 0
+    for i in range(0, x_all.shape[0], batch_size):
+        x = x_all[i:i + batch_size].to(device)
+        logit_s, logit_c = model(x)
+        pred_s = logit_s.argmax(dim=1)
+        pred_c = logit_c.argmax(dim=1)
+        hits += int(((pred_s == target_s) & (pred_c == target_c)).sum().item())
+    return hits / float(x_all.shape[0])

sprig/eval/prompts.py

@@ -0,0 +1,141 @@
+"""Fixed 32-prompt evaluation bank, minimal pairs, and held-out combos.
+
+The bank composition is pinned by the project plan (Part 3):
+8 seen single-object, 8 relations, 3 counts, 1 containment, 2 sizes,
+2 backgrounds, 4 held-out combos, 2 partial, 2 CLEVR-style = 32 prompts.
+
+Prompt phrasings follow the training caption templates in
+sprig/data/procgen/captions.py (T1/T2/T3/T5/T6/T7/T9 realizations; the
+attribute-dropped forms are valid partial captions), and the CLEVR-style
+prompts follow sprig/data/clevr/prep.py's synthesized caption grammar.
+
+The four compositional holdout combos {blue triangle, red ring, green star,
+yellow cross} never appear in training data; they appear here (only) as the
+held-out generalization probes.
+"""
+from __future__ import annotations
+
+from typing import List, Tuple
+
+# Canonical attribute vocabularies — single source of truth is
+# sprig/data/procgen/vocab.py; literals kept as fallback for lean imports.
+try:
+    from sprig.data.procgen.vocab import COLOR_NAMES as _COLOR_NAMES
+    from sprig.data.procgen.vocab import SHAPES as _SHAPES
+
+    COLORS: List[str] = list(_COLOR_NAMES)
+    SHAPES: List[str] = list(_SHAPES)
+except Exception:  # pragma: no cover
+    COLORS = ["red", "green", "blue", "yellow", "orange", "purple", "cyan", "magenta"]
+    SHAPES = ["circle", "square", "triangle", "rectangle", "diamond", "star", "cross", "ring"]
+
+SIZES: List[str] = ["small", "large"]
+RELATIONS: List[str] = ["to the left of", "to the right of", "above", "below"]
+
+# (color, shape) pairs excluded from all training scenes/captions.
+HELDOUT_COMBOS: List[Tuple[str, str]] = [
+    ("blue", "triangle"),
+    ("red", "ring"),
+    ("green", "star"),
+    ("yellow", "cross"),
+]
+
+_SINGLE_OBJECT = [
+    "a red circle",
+    "a blue square",
+    "a green triangle",
+    "a yellow star",
+    "a purple diamond",
+    "an orange cross",
+    "a cyan ring",
+    "a magenta rectangle",
+]
+
+_RELATIONS = [
+    "a red circle to the left of a blue square",
+    "a green diamond to the right of a yellow circle",
+    "a purple square above an orange circle",
+    "a cyan triangle below a magenta square",
+    "a yellow rectangle to the left of a purple star",
+    "an orange square to the right of a cyan diamond",
+    "a magenta circle above a red square",
+    "a blue diamond below a green rectangle",
+]
+
+_COUNTS = [
+    "a scene with two shapes on a white background",
+    "three shapes: a red circle, a blue square, and a green diamond",
+    "a scene with four shapes",
+]
+
+_CONTAINMENT = [
+    "a yellow circle inside a purple frame",
+]
+
+_SIZES = [
+    "a large orange diamond",
+    "a small cyan square",
+]
+
+_BACKGROUNDS = [
+    "a purple circle on a black background",
+    "a red square on a gray background",
+]
+
+_HELDOUT = ["a {} {}".format(c, s) for (c, s) in HELDOUT_COMBOS]
+
+_PARTIAL = [
+    "a circle",
+    "a striped square",
+]
+
+_CLEVR_STYLE = [
+    "a large red rubber cube to the left of a small blue metal sphere",
+    "a scene with three objects, including a small green metal cylinder",
+]
+
+PROMPTS: List[str] = (
+    _SINGLE_OBJECT
+    + _RELATIONS
+    + _COUNTS
+    + _CONTAINMENT
+    + _SIZES
+    + _BACKGROUNDS
+    + _HELDOUT
+    + _PARTIAL
+    + _CLEVR_STYLE
+)
+
+# Named groups for reporting / grid row labels.
+PROMPT_GROUPS = {
+    "single_object": _SINGLE_OBJECT,
+    "relations": _RELATIONS,
+    "counts": _COUNTS,
+    "containment": _CONTAINMENT,
+    "sizes": _SIZES,
+    "backgrounds": _BACKGROUNDS,
+    "heldout_combos": _HELDOUT,
+    "partial": _PARTIAL,
+    "clevr_style": _CLEVR_STYLE,
+}
+
+# (prompt_a, prompt_b, changed_attribute); attribute in
+# {"color", "shape", "relation", "size"}.
+MINIMAL_PAIRS: List[Tuple[str, str, str]] = [
+    ("a red circle", "a blue circle", "color"),
+    ("a green square", "a purple square", "color"),
+    ("a yellow diamond", "a cyan diamond", "color"),
+    ("an orange ring", "a magenta ring", "color"),
+    (
+        "a red circle to the left of a blue square",
+        "a red circle to the right of a blue square",
+        "relation",
+    ),
+    (
+        "a green circle above a yellow square",
+        "a green circle below a yellow square",
+        "relation",
+    ),
+    ("a small purple square", "a large purple square", "size"),
+    ("a red circle", "a red square", "shape"),
+]

sprig/eval/report.py

@@ -0,0 +1,657 @@
+"""Full evaluation report for a checkpoint: metrics.json + final_report.md.
+
+Expected data_dir layout (plan Part 1/3 formats):
+  data_dir/val/          images.u8 [N,64,64,3], emb.f16 (packed ragged),
+                         emb_offsets.i64 [N+1], meta.jsonl
+  data_dir/parse_eval/   same files; meta lines carry the GT tree under
+                         "tree" (or "gt_tree")
+  data_dir/train/images.u8   (optional; B0 is fit on val images otherwise)
+  data_dir/t5/null.f16       null-caption embedding, packed f16 [L0*768]
+  data_dir/t5/promptbank.npz npz {emb: f16 [P,Lmax,768] zero-padded,
+                         len: i32 [P]} from sprig.data.embed_t5 --prompts-out,
+                         32 prompts in sprig.eval.prompts.PROMPTS order
+  data_dir/t5/minimal_pairs.npz  same npz format, 16 prompts interleaved
+                         a0,b0,a1,b1,... in MINIMAL_PAIRS order (optional —
+                         the prompt-control gate is SKIPPED if absent)
+
+Stages that need a missing optional input are recorded as None and their
+gate reported SKIPPED, so the report degrades gracefully while any gate that
+can be evaluated is a hard PASS/FAIL.
+"""
+from __future__ import annotations
+
+import json
+import math
+import os
+from typing import Dict, List, Optional, Sequence, Tuple
+
+import numpy as np
+import torch
+from PIL import Image
+
+from sprig.eval import color_checks, monitors, tree_metrics
+from sprig.eval import baseline_pixmix
+from sprig.eval.prompts import COLORS, HELDOUT_COMBOS, MINIMAL_PAIRS, PROMPTS, SHAPES
+
+LOG2 = math.log(2.0)
+DIMS = 3.0 * 64 * 64
+
+GATE_THRESH = {
+    "b0_margin": 0.15,
+    "delta_c": 0.05,
+    "recall_tier1": 0.70,
+    "recall_tier2": 0.50,
+    "visible_cut_f1": 0.60,
+    "attribute_move": 0.80,
+    "relation_accuracy": 0.70,
+    "holdout_probe": 0.60,
+    "s_eff_frac": 0.25,
+    "alive_texel_frac": 0.50,
+}
+
+
+# --------------------------------------------------------------- data access
+
+class MemmapSplit:
+    """Reader for one dataset split in the plan's raw-memmap format."""
+
+    def __init__(self, split_dir: str):
+        self.dir = split_dir
+        img_path = os.path.join(split_dir, "images.u8")
+        n_bytes = os.path.getsize(img_path)
+        self.n = n_bytes // (64 * 64 * 3)
+        self.images = np.memmap(img_path, dtype=np.uint8, mode="r", shape=(self.n, 64, 64, 3))
+        off_path = os.path.join(split_dir, "emb_offsets.i64")
+        emb_path = os.path.join(split_dir, "emb.f16")
+        self.emb_offsets = np.fromfile(off_path, dtype=np.int64)
+        n_rows = int(self.emb_offsets[-1])
+        self.emb = np.memmap(emb_path, dtype=np.float16, mode="r", shape=(n_rows, 768))
+        self.meta: List[dict] = []
+        meta_path = os.path.join(split_dir, "meta.jsonl")
+        if os.path.exists(meta_path):
+            with open(meta_path) as f:
+                self.meta = [json.loads(line) for line in f if line.strip()]
+
+    def __len__(self) -> int:
+        return self.n
+
+    def image(self, i: int) -> np.ndarray:
+        return np.asarray(self.images[i])
+
+    def emb_i(self, i: int) -> torch.Tensor:
+        lo, hi = int(self.emb_offsets[i]), int(self.emb_offsets[i + 1])
+        return torch.from_numpy(np.array(self.emb[lo:hi])).to(torch.float16)
+
+    def tier(self, i: int) -> int:
+        return int(self.meta[i].get("tier", 0)) if self.meta else 0
+
+    def tree(self, i: int) -> Optional[dict]:
+        if not self.meta:
+            return None
+        return self.meta[i].get("tree") or self.meta[i].get("gt_tree")
+
+
+def load_prompt_npz(path: str) -> List[torch.Tensor]:
+    """embed_t5 --prompts-out npz {emb [P,Lmax,768] f16, len [P] i32} ->
+    list of unpadded [L_i,768] f16 tensors."""
+    z = np.load(path)
+    emb, lens = z["emb"], z["len"]
+    return [
+        torch.from_numpy(np.array(emb[i, : int(lens[i])])).to(torch.float16)
+        for i in range(emb.shape[0])
+    ]
+
+
+def load_null_emb(t5_dir: str) -> torch.Tensor:
+    flat = np.fromfile(os.path.join(t5_dir, "null.f16"), dtype=np.float16)
+    return torch.from_numpy(flat.reshape(-1, 768)).to(torch.float16)
+
+
+def pad_embs(embs: Sequence[torch.Tensor]) -> Tuple[torch.Tensor, torch.Tensor]:
+    """Ragged [L_i,768] f16 list -> (emb [B,Lmax,768] f16, emb_len [B] i32)."""
+    lens = torch.tensor([e.shape[0] for e in embs], dtype=torch.int32)
+    lmax = int(lens.max())
+    out = torch.zeros(len(embs), lmax, 768, dtype=torch.float16)
+    for i, e in enumerate(embs):
+        out[i, : e.shape[0]] = e
+    return out, lens
+
+
+# -------------------------------------------------------------- model access
+
+def _build_model_from_cfg(cfg):
+    """Instantiate SPRIGModel from a checkpoint 'config' entry, which may be a
+    SPRIGConfig, a full train-yaml dict (model section nested), or a flat dict
+    of SPRIGConfig fields."""
+    import dataclasses
+
+    from sprig.model.sprig import SPRIGConfig, SPRIGModel  # lazy: heavy import
+
+    if isinstance(cfg, SPRIGConfig):
+        return SPRIGModel(cfg)
+    if isinstance(cfg, dict):
+        section = cfg.get("model", cfg)
+        if isinstance(section, dict):
+            fields = {f.name for f in dataclasses.fields(SPRIGConfig)}
+            kw = {k: v for k, v in section.items() if k in fields}
+            return SPRIGModel(SPRIGConfig(**kw))
+    try:
+        return SPRIGModel(cfg)
+    except TypeError:
+        return SPRIGModel(**cfg)
+
+
+def load_model(ckpt_path: str, device: str = "cpu"):
+    """Load a SPRIGModel from a training checkpoint (lazy model import).
+
+    Prefers EMA weights (DESIGN: EMA is eval-only). train.py checkpoints
+    store EMA as {"decay", "shadow": {param_name: tensor}} — the shadow is
+    overlaid on the raw 'model' state dict (which also carries the buffers).
+    """
+    ckpt = torch.load(ckpt_path, map_location=device, weights_only=False)
+    if isinstance(ckpt, torch.nn.Module):
+        model = ckpt
+    else:
+        cfg = ckpt.get("config") or ckpt.get("cfg") or {}
+        model = _build_model_from_cfg(cfg)
+        state = None
+        for key in ("ema", "ema_state_dict", "model", "model_state_dict", "state_dict"):
+            if key in ckpt:
+                state = ckpt[key]
+                break
+        if state is None:
+            state = ckpt
+        if isinstance(state, dict) and isinstance(state.get("shadow"), dict):
+            base = dict(ckpt.get("model") or {})
+            base.update(state["shadow"])
+            state = base
+        model.load_state_dict(state)
+    model = model.to(device)
+    # Report numbers are untempered by definition: mid-anneal checkpoints carry
+    # a nonzero eta buffer which silently deflates bpd/delta_c (eval-audit
+    # finding C1 — the 75k ckpt reported bpd 0.53 instead of 3.97).
+    eta = getattr(model, "eta", None)
+    if isinstance(eta, torch.Tensor):
+        eta.data.fill_(0.0)
+    model.eval()
+    if hasattr(model, "report_mode"):
+        model.report_mode = True
+    return model
+
+
+def _as_u8_numpy(images) -> np.ndarray:
+    if isinstance(images, torch.Tensor):
+        images = images.detach().cpu().numpy()
+    return np.asarray(images).astype(np.uint8)
+
+
+def _log_marginal(model, images: np.ndarray, embs: Sequence[torch.Tensor], device: str) -> np.ndarray:
+    emb, emb_len = pad_embs(embs)
+    img = torch.from_numpy(np.array(images, dtype=np.uint8)).to(device)
+    with torch.no_grad():
+        logz = model.log_marginal(img, emb.to(device), emb_len.to(device))
+    return logz.detach().cpu().to(torch.float64).numpy()
+
+
+def bpd_from_logz(logz: np.ndarray) -> np.ndarray:
+    return -np.asarray(logz, dtype=np.float64) / (DIMS * LOG2)
+
+
+def _sample_images(model, emb: torch.Tensor, seed_struct: int, seed_material: int,
+                   n: int, device: str) -> np.ndarray:
+    e, elen = pad_embs([emb])
+    with torch.no_grad():
+        images, _trees = model.sample(
+            e.to(device), elen.to(device), seed_struct, seed_material, n
+        )
+    return _as_u8_numpy(images)
+
+
+def _parse_one(model, image: np.ndarray, emb: torch.Tensor, device: str):
+    e, elen = pad_embs([emb])
+    img = torch.from_numpy(np.array(image[None], dtype=np.uint8)).to(device)
+    with torch.no_grad():
+        result = model.map_parse(img, e.to(device), elen.to(device))
+    if result and isinstance(result[0], (list, tuple)):  # batched return
+        result = result[0]
+    return list(result)
+
+
+# ------------------------------------------------------------- image grids
+
+def image_grid(rows: Sequence[Sequence[np.ndarray]], pad: int = 2) -> Image.Image:
+    """[R][C] of u8 [64,64,3] arrays -> one PIL image with white padding."""
+    r, c = len(rows), max(len(row) for row in rows)
+    h = w = 64
+    canvas = np.full((r * (h + pad) + pad, c * (w + pad) + pad, 3), 255, dtype=np.uint8)
+    for i, row in enumerate(rows):
+        for j, img in enumerate(row):
+            y, x = pad + i * (h + pad), pad + j * (w + pad)
+            canvas[y:y + h, x:x + w] = np.asarray(img, dtype=np.uint8)
+    return Image.fromarray(canvas)
+
+
+def save_grid_jpeg(rows: Sequence[Sequence[np.ndarray]], path: str) -> None:
+    image_grid(rows).save(path, format="JPEG", quality=92)
+
+
+# ------------------------------------------------------------ eval stages
+
+def eval_bpd(model, split: MemmapSplit, device: str, max_images: int = 2000,
+             batch_size: int = 32) -> Dict[str, object]:
+    n = min(len(split), max_images)
+    bpds = np.zeros(n)
+    tiers = np.array([split.tier(i) for i in range(n)], dtype=np.int64)
+    for i in range(0, n, batch_size):
+        j = min(i + batch_size, n)
+        embs = [split.emb_i(k) for k in range(i, j)]
+        logz = _log_marginal(model, np.asarray(split.images[i:j]), embs, device)
+        bpds[i:j] = bpd_from_logz(logz)
+    per_tier = {
+        int(t): float(bpds[tiers == t].mean()) for t in np.unique(tiers)
+    }
+    ge1 = bpds[tiers >= 1]
+    return {
+        "bpd_val": float(bpds.mean()),
+        "bpd_per_tier": per_tier,
+        "bpd_tier_ge1": float(ge1.mean()) if ge1.size else float(bpds.mean()),
+        "tier_ge1_index": np.nonzero(tiers >= 1)[0],
+        "n": n,
+    }
+
+
+def eval_delta_c(model, split: MemmapSplit, null_emb: torch.Tensor, device: str,
+                 max_images: int = 512, batch_size: int = 32) -> float:
+    n = min(len(split), max_images)
+    total = 0.0
+    for i in range(0, n, batch_size):
+        j = min(i + batch_size, n)
+        imgs = np.asarray(split.images[i:j])
+        embs = [split.emb_i(k) for k in range(i, j)]
+        logz_c = _log_marginal(model, imgs, embs, device)
+        logz_0 = _log_marginal(model, imgs, [null_emb] * (j - i), device)
+        total += float((bpd_from_logz(logz_0) - bpd_from_logz(logz_c)).sum())
+    return total / n
+
+
+def eval_caption_swap(model, split: MemmapSplit, device: str, n_groups: int = 4,
+                      group: int = 8) -> float:
+    """8-way in-batch caption swap: fraction of images whose own caption gives
+    a higher logZ than all 7 swapped captions."""
+    rng = np.random.default_rng(0)
+    n = len(split)
+    wins = 0
+    total = 0
+    for _ in range(n_groups):
+        idx = rng.choice(n, size=group, replace=False)
+        imgs = np.asarray(split.images[np.sort(idx)])
+        embs = [split.emb_i(int(i)) for i in np.sort(idx)]
+        scores = np.zeros((group, group))
+        for j in range(group):
+            scores[:, j] = _log_marginal(model, imgs, [embs[j]] * group, device)
+        for i in range(group):
+            off_diag = np.delete(scores[i], i)
+            wins += int(scores[i, i] > off_diag.max())
+            total += 1
+    return wins / float(total)
+
+
+def eval_b0(fit_images: np.ndarray, eval_images: np.ndarray, steps: int = 2000,
+            device: str = "cpu") -> float:
+    b0 = baseline_pixmix.fit(fit_images, steps=steps, device=device)
+    return float(b0.bpd(eval_images))
+
+
+def eval_tree_metrics(model, split: MemmapSplit, device: str,
+                      max_images: int = 512) -> Dict[str, object]:
+    n = min(len(split), max_images)
+    recalls: Dict[int, List[float]] = {}
+    f1s: List[float] = []
+    aris: List[float] = []
+    for i in range(n):
+        gt = split.tree(i)
+        if gt is None:
+            continue
+        img = split.image(i)
+        parse = _parse_one(model, img, split.emb_i(i), device)
+        recalls.setdefault(split.tier(i), []).append(
+            tree_metrics.object_cell_recall(parse, gt)
+        )
+        f1s.append(tree_metrics.visible_cut_f1(parse, gt, img))
+        aris.append(tree_metrics.leaf_ari(parse, gt))
+    out: Dict[str, object] = {
+        "visible_cut_f1": float(np.mean(f1s)) if f1s else None,
+        "leaf_ari": float(np.mean(aris)) if aris else None,
+        "recall_per_tier": {t: float(np.mean(v)) for t, v in recalls.items()},
+    }
+    out["recall_tier1"] = out["recall_per_tier"].get(1)  # type: ignore[union-attr]
+    out["recall_tier2"] = out["recall_per_tier"].get(2)  # type: ignore[union-attr]
+    return out
+
+
+def sample_prompt_bank_grid(model, prompt_embs: Sequence[torch.Tensor], out_path: str,
+                            device: str, n_seeds: int = 8, seed0: int = 0) -> None:
+    rows = []
+    for i, emb in enumerate(prompt_embs):
+        row = [
+            _sample_images(model, emb, seed0 + 1000 * s + i, seed0 + 5000 + 1000 * s + i,
+                           1, device)[0]
+            for s in range(n_seeds)
+        ]
+        rows.append(row)
+    save_grid_jpeg(rows, out_path)
+
+
+def layout_material_grid(model, emb: torch.Tensor, out_path: str, device: str,
+                         k: int = 4) -> None:
+    """k x k grid: rows = frozen structural seed, cols = rerolled material seed."""
+    rows = []
+    for r in range(k):
+        rows.append([
+            _sample_images(model, emb, 100 + r, 900 + c, 1, device)[0] for c in range(k)
+        ])
+    save_grid_jpeg(rows, out_path)
+
+
+def _tokens(prompt: str) -> List[str]:
+    return prompt.replace(",", " ").split()
+
+def _color_in(prompt: str) -> List[str]:
+    return [t for t in _tokens(prompt) if t in COLORS]
+
+def _shape_in(prompt: str) -> List[str]:
+    return [t for t in _tokens(prompt) if t in SHAPES]
+
+def _relation_in(prompt: str) -> Optional[str]:
+    for rel in ("left of", "right of"):
+        if rel in prompt:
+            return rel
+    for rel in ("above", "below"):
+        if rel in prompt.split():
+            return rel
+    return None
+
+
+def _median_object_area(images: np.ndarray) -> float:
+    areas = []
+    for img in images:
+        objs = color_checks.extract(img)["objects"]
+        if objs:
+            areas.append(objs[0]["area"])
+    return float(np.median(areas)) if areas else 0.0
+
+
+def eval_minimal_pairs(model, pair_embs: Sequence[torch.Tensor], device: str,
+                       n_seeds: int = 64, probe_ckpt: Optional[str] = None,
+                       batch: int = 16) -> Dict[str, object]:
+    """pair_embs: 16 embeddings, (a_i, b_i) interleaved, MINIMAL_PAIRS order."""
+    per_pair: Dict[str, float] = {}
+    move_scores: List[float] = []
+    rel_scores: List[float] = []
+    for p, (prompt_a, prompt_b, attr) in enumerate(MINIMAL_PAIRS):
+        emb_a, emb_b = pair_embs[2 * p], pair_embs[2 * p + 1]
+        imgs_b = np.concatenate([
+            _sample_images(model, emb_b, 10 * p + s, 77 + 10 * p + s, 1, device)
+            for s in range(n_seeds)
+        ])
+        score: Optional[float] = None
+        if attr == "color":
+            new_color = [c for c in _color_in(prompt_b) if c not in _color_in(prompt_a)][0]
+            hits = [
+                any(o["color"] == new_color for o in color_checks.extract(im)["objects"])
+                for im in imgs_b
+            ]
+            score = float(np.mean(hits))
+        elif attr == "relation":
+            colors = _color_in(prompt_b)
+            rel = _relation_in(prompt_b)
+            oks = []
+            for im in imgs_b:
+                r = color_checks.relation_holds(
+                    color_checks.extract(im), colors[0], colors[1], rel
+                )
+                oks.append(bool(r) if r is not None else False)
+            score = float(np.mean(oks))
+        elif attr == "size":
+            imgs_a = np.concatenate([
+                _sample_images(model, emb_a, 10 * p + s, 77 + 10 * p + s, 1, device)
+                for s in range(n_seeds)
+            ])
+            bigger_in_b = "large" in _tokens(prompt_b)
+            area_a, area_b = _median_object_area(imgs_a), _median_object_area(imgs_b)
+            score = float(area_b > area_a) if bigger_in_b else float(area_b < area_a)
+        elif attr == "shape" and probe_ckpt is not None:
+            from sprig.eval import probe
+
+            score = probe.score_generations(
+                imgs_b, _shape_in(prompt_b)[0], _color_in(prompt_b)[0], probe_ckpt, device
+            )
+        if score is not None:
+            per_pair["{} -> {}".format(prompt_a, prompt_b)] = score
+            (rel_scores if attr == "relation" else move_scores).append(score)
+    return {
+        "per_pair": per_pair,
+        "attribute_move": float(np.mean(move_scores)) if move_scores else None,
+        "relation_accuracy": float(np.mean(rel_scores)) if rel_scores else None,
+    }
+
+
+def eval_heldout_combos(model, prompt_embs: Sequence[torch.Tensor], probe_ckpt: str,
+                        device: str, n: int = 64) -> Dict[str, object]:
+    from sprig.eval import probe
+
+    model_probe = probe.load_probe(probe_ckpt, device)
+    per_combo: Dict[str, float] = {}
+    for color, shape in HELDOUT_COMBOS:
+        prompt = "a {} {}".format(color, shape)
+        emb = prompt_embs[PROMPTS.index(prompt)]
+        imgs = np.concatenate([
+            _sample_images(model, emb, 313 + s, 707 + s, 1, device) for s in range(n)
+        ])
+        per_combo[prompt] = probe.score_generations(imgs, shape, color, model_probe, device)
+    return {
+        "per_combo": per_combo,
+        "holdout_probe_acc": float(np.mean(list(per_combo.values()))),
+    }
+
+
+def eval_health(model, split: MemmapSplit, device: str, n_images: int = 64) -> Dict[str, object]:
+    n = min(len(split), n_images)
+    imgs = torch.from_numpy(np.array(split.images[:n], dtype=np.uint8)).to(device)
+    emb, elen = pad_embs([split.emb_i(i) for i in range(n)])
+    stats = model.posterior_usage(imgs, emb.to(device), elen.to(device))
+    usage = stats["symbol_usage"]
+    texels = stats["texel_usage"]
+    s = int(usage.shape[-1]) if hasattr(usage, "shape") else len(usage)
+    t_v = int(texels.shape[-1]) if hasattr(texels, "shape") else len(texels)
+    return {
+        "S": s,
+        "T_v": t_v,
+        "s_eff": monitors.s_eff(usage),
+        "alive_texel_frac": monitors.alive_texels(texels, t_v) / float(t_v),
+        "node_entropy": float(stats.get("node_entropy", float("nan"))),
+    }
+
+
+# ---------------------------------------------------------------- gates
+
+def _gate(ok: Optional[bool], detail: str) -> Dict[str, str]:
+    if ok is None:
+        return {"status": "SKIPPED", "detail": detail}
+    return {"status": "PASS" if ok else "FAIL", "detail": detail}
+
+
+def evaluate_gates(metrics: Dict[str, object]) -> Dict[str, Dict[str, str]]:
+    """The plan's 5 proof-of-concept success criteria as PASS/FAIL/SKIPPED."""
+    g = GATE_THRESH
+    gates: Dict[str, Dict[str, str]] = {}
+
+    bpd = metrics.get("bpd_tier_ge1")
+    b0 = metrics.get("b0_bpd")
+    dc = metrics.get("delta_c")
+    if bpd is None or b0 is None or dc is None:
+        gates["1_likelihood"] = _gate(None, "missing bpd/b0/delta_c")
+    else:
+        ok = (b0 - bpd >= g["b0_margin"]) and (dc >= g["delta_c"])
+        gates["1_likelihood"] = _gate(ok, "bpd(tier>=1)={:.3f} vs B0={:.3f} (need margin >= {}); delta_c={:.3f} (need >= {})".format(bpd, b0, g["b0_margin"], dc, g["delta_c"]))
+
+    tm = metrics.get("tree") or {}
+    r1, r2, f1 = tm.get("recall_tier1"), tm.get("recall_tier2"), tm.get("visible_cut_f1")
+    if r1 is None or r2 is None or f1 is None:
+        gates["2_parses"] = _gate(None, "missing tree metrics")
+    else:
+        ok = (r1 >= g["recall_tier1"]) and (r2 >= g["recall_tier2"]) and (f1 >= g["visible_cut_f1"])
+        gates["2_parses"] = _gate(ok, "recall t1={:.2f} (>= {}), t2={:.2f} (>= {}), cut F1={:.2f} (>= {})".format(r1, g["recall_tier1"], r2, g["recall_tier2"], f1, g["visible_cut_f1"]))
+
+    ps = metrics.get("prompt_swap") or {}
+    move, rel = ps.get("attribute_move"), ps.get("relation_accuracy")
+    if move is None or rel is None:
+        gates["3_prompt_control"] = _gate(None, "missing minimal-pair scores")
+    else:
+        ok = (move >= g["attribute_move"]) and (rel >= g["relation_accuracy"])
+        gates["3_prompt_control"] = _gate(ok, "attribute-move={:.2f} (>= {}), relation={:.2f} (>= {})".format(move, g["attribute_move"], rel, g["relation_accuracy"]))
+
+    hp = metrics.get("holdout_probe_acc")
+    if hp is None:
+        gates["4_compositional"] = _gate(None, "no probe checkpoint / heldout scores")
+    else:
+        gates["4_compositional"] = _gate(hp >= g["holdout_probe"], "held-out combo probe acc={:.2f} (>= {})".format(hp, g["holdout_probe"]))
+
+    health = metrics.get("health") or {}
+    s_eff, s = health.get("s_eff"), health.get("S")
+    alive = health.get("alive_texel_frac")
+    if s_eff is None or alive is None or s is None:
+        gates["5_health"] = _gate(None, "missing health stats")
+    else:
+        ok = (s_eff >= g["s_eff_frac"] * s) and (alive >= g["alive_texel_frac"])
+        gates["5_health"] = _gate(ok, "S_eff={:.0f} (>= {:.0f}), alive texels={:.0%} (>= {:.0%})".format(s_eff, g["s_eff_frac"] * s, alive, g["alive_texel_frac"]))
+    return gates
+
+
+def write_report(metrics: Dict[str, object], gates: Dict[str, Dict[str, str]],
+                 out_dir: str) -> None:
+    os.makedirs(out_dir, exist_ok=True)
+
+    def _default(o):
+        if isinstance(o, (np.integer,)):
+            return int(o)
+        if isinstance(o, (np.floating,)):
+            return float(o)
+        if isinstance(o, np.ndarray):
+            return o.tolist()
+        return str(o)
+
+    with open(os.path.join(out_dir, "metrics.json"), "w") as f:
+        json.dump({"metrics": metrics, "gates": gates}, f, indent=2, default=_default)
+    lines = ["# SPRIG v0.1 — Final Report", "", "## Success-criteria gates", ""]
+    lines.append("| gate | status | detail |")
+    lines.append("|---|---|---|")
+    for name, g in gates.items():
+        lines.append("| {} | **{}** | {} |".format(name, g["status"], g["detail"]))
+    lines += ["", "## Key metrics", "", "```json"]
+    lines.append(json.dumps(metrics, indent=2, default=_default))
+    lines.append("```")
+    with open(os.path.join(out_dir, "final_report.md"), "w") as f:
+        f.write("\n".join(lines) + "\n")
+
+
+# ------------------------------------------------------------- orchestrator
+
+def run_report(
+    ckpt_path: Optional[str],
+    data_dir: str,
+    out_dir: str,
+    device: str = "cpu",
+    model=None,
+    probe_ckpt: Optional[str] = None,
+    n_bank_seeds: int = 8,
+    n_pair_seeds: int = 64,
+    max_bpd_images: int = 2000,
+    max_parse_images: int = 512,
+    b0_steps: int = 2000,
+) -> Dict[str, object]:
+    """Run the full evaluation suite on a checkpoint.
+
+    Pass `model` directly to skip checkpoint loading (used by tests/harness).
+    Returns the metrics dict; writes metrics.json, final_report.md and the
+    sample grids into out_dir.
+    """
+    os.makedirs(out_dir, exist_ok=True)
+    if model is None:
+        model = load_model(ckpt_path, device)
+
+    val_dir = os.path.join(data_dir, "val")
+    if not os.path.isdir(val_dir):
+        val_dir = os.path.join(data_dir, "val_fast")
+    val = MemmapSplit(val_dir)
+    t5_dir = os.path.join(data_dir, "t5")
+
+    metrics: Dict[str, object] = {}
+
+    # --- likelihood
+    bpd_stats = eval_bpd(model, val, device, max_images=max_bpd_images)
+    tier_ge1_index = bpd_stats.pop("tier_ge1_index")
+    metrics.update(bpd_stats)
+
+    null_path = os.path.join(t5_dir, "null.f16")
+    metrics["delta_c"] = (
+        eval_delta_c(model, val, load_null_emb(t5_dir), device)
+        if os.path.exists(null_path) else None
+    )
+    metrics["caption_swap_win_frac"] = eval_caption_swap(model, val, device)
+
+    # --- B0 baseline (fit on train images when available, else val)
+    train_imgs_path = os.path.join(data_dir, "train", "images.u8")
+    if os.path.exists(train_imgs_path):
+        n_fit = os.path.getsize(train_imgs_path) // (64 * 64 * 3)
+        fit_imgs = np.memmap(train_imgs_path, dtype=np.uint8, mode="r",
+                             shape=(n_fit, 64, 64, 3))[: 4096]
+    else:
+        fit_imgs = np.asarray(val.images[: min(len(val), 4096)])
+    eval_idx = tier_ge1_index if len(tier_ge1_index) else np.arange(min(len(val), 512))
+    eval_imgs = np.asarray(val.images)[eval_idx[:512]]
+    metrics["b0_bpd"] = eval_b0(np.asarray(fit_imgs), eval_imgs, steps=b0_steps, device=device)
+
+    # --- parse metrics
+    parse_dir = os.path.join(data_dir, "parse_eval")
+    metrics["tree"] = (
+        eval_tree_metrics(model, MemmapSplit(parse_dir), device, max_parse_images)
+        if os.path.isdir(parse_dir) else None
+    )
+
+    # --- prompt bank grids
+    bank_npz = os.path.join(t5_dir, "promptbank.npz")
+    prompt_embs: Optional[List[torch.Tensor]] = None
+    if os.path.exists(bank_npz):
+        prompt_embs = load_prompt_npz(bank_npz)
+        sample_prompt_bank_grid(
+            model, prompt_embs, os.path.join(out_dir, "prompt_bank_grid.jpg"),
+            device, n_seeds=n_bank_seeds,
+        )
+        layout_material_grid(
+            model, prompt_embs[8], os.path.join(out_dir, "layout_material_grid.jpg"), device
+        )
+
+    # --- minimal pairs
+    pairs_npz = os.path.join(t5_dir, "minimal_pairs.npz")
+    if os.path.exists(pairs_npz):
+        pair_embs = load_prompt_npz(pairs_npz)
+        metrics["prompt_swap"] = eval_minimal_pairs(
+            model, pair_embs, device, n_seeds=n_pair_seeds, probe_ckpt=probe_ckpt
+        )
+    else:
+        metrics["prompt_swap"] = None
+
+    # --- held-out combos via probe
+    metrics["holdout_probe_acc"] = None
+    if probe_ckpt is not None and prompt_embs is not None:
+        heldout = eval_heldout_combos(model, prompt_embs, probe_ckpt, device)
+        metrics["heldout_per_combo"] = heldout["per_combo"]
+        metrics["holdout_probe_acc"] = heldout["holdout_probe_acc"]
+
+    # --- health
+    metrics["health"] = eval_health(model, val, device)
+
+    gates = evaluate_gates(metrics)
+    write_report(metrics, gates, out_dir)
+    return metrics

sprig/eval/tree_metrics.py

@@ -0,0 +1,230 @@
+"""Parse metrics vs ground-truth region trees (identifiability-aware).
+
+Node schema (duck-typed; both are accepted everywhere a tree is expected):
+- GT JSON node (meta.jsonl "tree", written by sprig/data/procgen/sampler.py):
+  internal {"rect": [x0,y0,x1,y1] px (x1/y1 exclusive), "axis": "V"|"H",
+  "cut": px, "children": [lo, hi]}; leaf {"rect", "leaf": true,
+  "obj": int|null, "fill": ...} — a leaf is an object-role leaf iff "obj"
+  is not null (an inline "object" dict or "role" == "object" also works).
+  "axis"/"cut" fields are ignored: cut geometry is derived from the children
+  rects, which makes the metrics robust to axis naming conventions.
+- Model ParseNode (sprig/model/sprig.py): attributes rect, children (list or
+  None), axis, cut_px, symbol, texel.
+
+`parse` arguments may be a single root node or the list returned by
+`model.map_parse` (all nodes are collected either way).
+"""
+from __future__ import annotations
+
+from typing import Dict, List, Optional, Sequence, Tuple
+
+import numpy as np
+
+Rect = Tuple[int, int, int, int]
+CutSeg = Tuple[str, float, float, float]  # (axis "V"|"H", position, ext_lo, ext_hi)
+
+
+# ---------------------------------------------------------------- accessors
+
+def _rect(node) -> Rect:
+    r = node["rect"] if isinstance(node, dict) else node.rect
+    x0, y0, x1, y1 = (int(v) for v in r)
+    return (x0, y0, x1, y1)
+
+
+def _children(node) -> List:
+    if isinstance(node, dict):
+        ch = node.get("children")
+    else:
+        ch = getattr(node, "children", None)
+    return list(ch) if ch else []
+
+
+def _object(node) -> Optional[dict]:
+    """Object payload of an object-role leaf, else None.
+
+    Accepts the sampler schema (leaf {"obj": <int index>|null}), an inline
+    {"object": {...}} dict, or a "role": "object" marker."""
+    if isinstance(node, dict):
+        if node.get("object") is not None:
+            return node["object"]
+        if node.get("obj") is not None:
+            return {"index": node["obj"]}
+        if node.get("role") == "object":
+            return {}
+        return None
+    return getattr(node, "object", None)
+
+
+def collect_nodes(tree_or_nodes) -> List:
+    """Flatten a root node / list of nodes into a deduped list of all nodes."""
+    roots = tree_or_nodes if isinstance(tree_or_nodes, (list, tuple)) else [tree_or_nodes]
+    out: List = []
+    seen = set()
+    stack = list(roots)
+    while stack:
+        n = stack.pop()
+        if id(n) in seen:
+            continue
+        seen.add(id(n))
+        out.append(n)
+        stack.extend(_children(n))
+    return out
+
+
+def leaves(tree_or_nodes) -> List:
+    return [n for n in collect_nodes(tree_or_nodes) if not _children(n)]
+
+
+def cut_segments(tree_or_nodes) -> List[CutSeg]:
+    """All cut segments, derived from each internal node's children rects."""
+    segs: List[CutSeg] = []
+    for n in collect_nodes(tree_or_nodes):
+        ch = _children(n)
+        if len(ch) != 2:
+            continue
+        (ax0, ay0, ax1, ay1), (bx0, by0, bx1, by1) = _rect(ch[0]), _rect(ch[1])
+        if ax1 == bx0 and ay0 == by0 and ay1 == by1:  # vertical line at x=ax1
+            segs.append(("V", float(ax1), float(ay0), float(ay1)))
+        elif bx1 == ax0 and ay0 == by0 and ay1 == by1:
+            segs.append(("V", float(bx1), float(ay0), float(ay1)))
+        elif ay1 == by0 and ax0 == bx0 and ax1 == bx1:  # horizontal line at y=ay1
+            segs.append(("H", float(ay1), float(ax0), float(ax1)))
+        elif by1 == ay0 and ax0 == bx0 and ax1 == bx1:
+            segs.append(("H", float(by1), float(ax0), float(ax1)))
+    return segs
+
+
+# ------------------------------------------------------------------ metrics
+
+def _iou(a: Rect, b: Rect) -> float:
+    ix = max(0, min(a[2], b[2]) - max(a[0], b[0]))
+    iy = max(0, min(a[3], b[3]) - max(a[1], b[1]))
+    inter = ix * iy
+    if inter == 0:
+        return 0.0
+    area_a = (a[2] - a[0]) * (a[3] - a[1])
+    area_b = (b[2] - b[0]) * (b[3] - b[1])
+    return inter / float(area_a + area_b - inter)
+
+
+def object_cell_recall(parse, gt, iou_thresh: float = 0.8) -> float:
+    """Fraction of GT object-role leaves whose rect is matched (IoU >= thresh)
+    by ANY node rect of the parse (internal or leaf). Vacuously 1.0 when the
+    GT tree has no object leaves."""
+    gt_obj_rects = [_rect(n) for n in leaves(gt) if _object(n) is not None]
+    if not gt_obj_rects:
+        return 1.0
+    parse_rects = [_rect(n) for n in collect_nodes(parse)]
+    hit = 0
+    for gr in gt_obj_rects:
+        if any(_iou(gr, pr) >= iou_thresh for pr in parse_rects):
+            hit += 1
+    return hit / float(len(gt_obj_rects))
+
+
+def _cut_contrast(image: np.ndarray, seg: CutSeg, band: int = 2) -> float:
+    """Max-channel abs difference of mean colors on the two sides of a cut."""
+    img = np.asarray(image, dtype=np.float64)
+    axis, pos, lo, hi = seg
+    p, a, b = int(round(pos)), int(round(lo)), int(round(hi))
+    if axis == "V":
+        left = img[a:b, max(0, p - band):p, :]
+        right = img[a:b, p:min(img.shape[1], p + band), :]
+    else:
+        left = img[max(0, p - band):p, a:b, :]
+        right = img[p:min(img.shape[0], p + band), a:b, :]
+    if left.size == 0 or right.size == 0:
+        return 0.0
+    return float(
+        np.abs(left.reshape(-1, 3).mean(axis=0) - right.reshape(-1, 3).mean(axis=0)).max()
+    )
+
+
+def _visible(segs: Sequence[CutSeg], image: np.ndarray, contrast_thresh: float) -> List[CutSeg]:
+    return [s for s in segs if _cut_contrast(image, s) > contrast_thresh]
+
+
+def _seg_match(a: CutSeg, b: CutSeg, tol_px: float) -> bool:
+    if a[0] != b[0] or abs(a[1] - b[1]) > tol_px:
+        return False
+    overlap = min(a[3], b[3]) - max(a[2], b[2])
+    shorter = min(a[3] - a[2], b[3] - b[2])
+    return shorter > 0 and overlap >= 0.5 * shorter
+
+
+def visible_cut_f1(
+    parse,
+    gt,
+    image: np.ndarray,
+    contrast_thresh: float = 20.0,
+    tol_px: float = 1.5,
+) -> float:
+    """Boundary F1 over VISIBLE cut segments.
+
+    Both GT and parse cut segments are filtered by actual cross-cut mean-color
+    contrast in `image` (> contrast_thresh, max over RGB channels); visible
+    parse segments are greedily one-to-one matched to visible GT segments
+    (same axis, position within tol_px, >=50% extent overlap).
+    """
+    gt_segs = _visible(cut_segments(gt), image, contrast_thresh)
+    pr_segs = _visible(cut_segments(parse), image, contrast_thresh)
+    if not gt_segs and not pr_segs:
+        return 1.0
+    if not gt_segs or not pr_segs:
+        return 0.0
+    used = [False] * len(gt_segs)
+    tp = 0
+    for ps in pr_segs:
+        cands = [
+            (abs(ps[1] - gs[1]), j)
+            for j, gs in enumerate(gt_segs)
+            if not used[j] and _seg_match(ps, gs, tol_px)
+        ]
+        if cands:
+            used[min(cands)[1]] = True
+            tp += 1
+    prec = tp / float(len(pr_segs))
+    rec = tp / float(len(gt_segs))
+    return 0.0 if tp == 0 else 2 * prec * rec / (prec + rec)
+
+
+def leaf_assignment_map(tree_or_nodes, canvas: int = 64) -> np.ndarray:
+    """[canvas,canvas] int32 map: pixel -> index of the leaf covering it."""
+    out = np.full((canvas, canvas), -1, dtype=np.int32)
+    for i, leaf in enumerate(leaves(tree_or_nodes)):
+        x0, y0, x1, y1 = _rect(leaf)
+        out[y0:y1, x0:x1] = i
+    return out
+
+
+def adjusted_rand_index(a: np.ndarray, b: np.ndarray) -> float:
+    """ARI between two integer label maps of identical shape (numpy only)."""
+    a = np.asarray(a).ravel()
+    b = np.asarray(b).ravel()
+    _, ai = np.unique(a, return_inverse=True)
+    _, bi = np.unique(b, return_inverse=True)
+    na, nb = ai.max() + 1, bi.max() + 1
+    cont = np.bincount(ai * nb + bi, minlength=na * nb).reshape(na, nb).astype(np.float64)
+    n = cont.sum()
+
+    def _comb2(x: np.ndarray) -> np.ndarray:
+        return x * (x - 1) / 2.0
+
+    sum_ij = _comb2(cont).sum()
+    sum_a = _comb2(cont.sum(axis=1)).sum()
+    sum_b = _comb2(cont.sum(axis=0)).sum()
+    total = _comb2(np.array(n))
+    expected = sum_a * sum_b / total if total > 0 else 0.0
+    max_index = 0.5 * (sum_a + sum_b)
+    denom = max_index - expected
+    if denom == 0:
+        return 1.0 if sum_ij == max_index else 0.0
+    return float((sum_ij - expected) / denom)
+
+
+def leaf_ari(parse, gt, canvas: int = 64) -> float:
+    """Adjusted Rand index between parse and GT pixel->leaf assignment maps."""
+    return adjusted_rand_index(
+        leaf_assignment_map(parse, canvas), leaf_assignment_map(gt, canvas)
+    )

sprig/model/atlas.py

@@ -0,0 +1,291 @@
+"""Texel atlas renderer and leaf emission scoring (DESIGN.md section 4, atlas.py).
+
+The renderer amortizes the terminal decoder into a canonical per-texel atlas of
+discretized-logistic parameters: atlas [B, T_v, 40, 16, 16] (channel layout in
+sprig/model/dl.py). A trainable per-texel additive bias grid [T_v, 40, 16, 16]
+is added to the renderer output — the resurrection-writable parameterization
+(F3.3 / M1.2): dead texels are revived by directly overwriting their bias rows.
+
+FiLM mapping (illumination field Phi [B, 8, 16, 16] over the canvas), sampled
+bilinearly at each leaf's canvas-center position -> phi [..., 8]:
+
+    scale_c = 1 + phi[c]      for c in {0: R, 1: G, 2: B}
+    shift_c =     phi[3 + c]  for c in {0: R, 1: G, 2: B}
+    phi[6:8] reserved (unused in v0.1)
+
+Each DL *mean* channel group (R/G/B means, identically across the 4 mixture
+components) is modulated as mean' = mean * scale_c + shift_c. Log-scales,
+weights and coupling coefficients are not modulated.
+"""
+from __future__ import annotations
+
+import math
+import os
+from typing import Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sprig.model import dl
+from sprig.model.gmt import CrossAttnBlock, caption_padding_mask
+
+ATLAS_RES = 16
+ATLAS_WIDTH = 256
+
+
+_LOG_HALF_RANGE = math.log(127.5)
+
+
+def _ell_chunk_math(
+    pooled_g: torch.Tensor,
+    crops_c: torch.Tensor,
+    w_c: torch.Tensor,
+    sc_c: torch.Tensor,
+    sh_c: torch.Tensor,
+) -> torch.Tensor:
+    """One region-chunk of emission scores -> [B, ng, T_v] fp32.
+
+    Runs under gradient checkpointing: the per-pixel intermediates are
+    recomputed in backward instead of stored (storing them across all chunks
+    is >90GB at batch 256). Exactly the dl.dl_logprob math with FiLM'd means
+    (fp32 throughout), but unrolled over the 4 mixture components x 3 RGB
+    channels so every intermediate is a [B,ng,T_v,h,w] slab instead of the
+    [B,ng,T_v,4,3,h,w] block — torch.compile then fuses the whole chunk
+    without materializing multi-GiB buffers (the tensorized form made
+    inductor allocate the full block, ~18 GiB at B=128 leaf_chunk=49).
+    """
+    # NOTE: a bf16 variant (inputs cast down, fp32 pixel-sum) was tried and
+    # REVERTED: eager bf16 shifts ell by tens of nats (bf16 ulp at |ell|~1e3
+    # is ~8) and fails the 0.1% loss parity gate by ~40x.
+    p = pooled_g.float()                                     # [B,T_v,40,h,w]
+    x = crops_c.float()                                      # [B,ng,3,h,w]
+    sc = sc_c.float()                                        # [B,ng,3]
+    sh = sh_c.float()
+    # Mixture-weight log-softmax over the 4 weight channels (pooled-sized).
+    wl = [p[:, :, 10 * j] for j in range(dl.N_COMP)]         # [B,T_v,h,w] each
+    wm = torch.logaddexp(torch.logaddexp(wl[0], wl[1]),
+                         torch.logaddexp(wl[2], wl[3]))      # logsumexp_j
+    x_ch = [x[:, :, c].unsqueeze(2) for c in range(3)]       # [B,ng,1,h,w]
+
+    mix: Optional[torch.Tensor] = None
+    for j in range(dl.N_COMP):
+        acc: Optional[torch.Tensor] = None                   # sum_c log P_j(c)
+        for c in range(3):
+            mu = p[:, :, 10 * j + 1 + c].unsqueeze(1)        # [B,1,T_v,h,w]
+            ls = p[:, :, 10 * j + 4 + c].clamp(
+                dl.LOGSCALE_MIN, dl.LOGSCALE_MAX).unsqueeze(1)
+            m = mu * sc[:, :, c, None, None, None] + sh[:, :, c, None, None, None]
+            if c == 1:
+                al = torch.tanh(p[:, :, 10 * j + 7]).unsqueeze(1)
+                m = m + al * x_ch[0]
+            elif c == 2:
+                be = torch.tanh(p[:, :, 10 * j + 8]).unsqueeze(1)
+                ga = torch.tanh(p[:, :, 10 * j + 9]).unsqueeze(1)
+                m = m + be * x_ch[0] + ga * x_ch[1]
+            xc = x_ch[c]
+            centered = xc - m                                # [B,ng,T_v,h,w]
+            inv_std = torch.exp(-ls)
+            plus_in = inv_std * (centered + dl._HALF_BIN)
+            min_in = inv_std * (centered - dl._HALF_BIN)
+            log_cdf_plus = plus_in - F.softplus(plus_in)
+            log_one_minus_cdf_min = -F.softplus(min_in)
+            cdf_delta = torch.sigmoid(plus_in) - torch.sigmoid(min_in)
+            mid_in = inv_std * centered
+            log_pdf_mid = mid_in - ls - 2.0 * F.softplus(mid_in)
+            log_prob_mid = torch.where(
+                cdf_delta > 1e-5,
+                torch.log(cdf_delta.clamp(min=1e-12)),
+                log_pdf_mid - _LOG_HALF_RANGE,
+            )
+            lp = torch.where(
+                xc < -0.999,
+                log_cdf_plus,
+                torch.where(xc > 0.999, log_one_minus_cdf_min, log_prob_mid),
+            )
+            acc = lp if acc is None else acc + lp
+        pc = acc + (wl[j] - wm).unsqueeze(1)                 # log w_j + log P_j
+        mix = pc if mix is None else torch.logaddexp(mix, pc)
+    # Per-pixel importance weights (object-pixel up-weighting; all-ones in the
+    # unweighted/eval path — multiplying by exactly 1.0 is bit-identical fp32,
+    # so likelihood parity is preserved).
+    return (mix * w_c.float().unsqueeze(2)).sum(dim=(-1, -2)).float()  # [B,ng,T_v]
+
+
+_ELL_COMPILED = None  # lazily-built torch.compile wrapper of _ell_chunk_math
+
+
+def _ell_chunk_fn(device: torch.device):
+    """Fused (torch.compile) emission-chunk kernel on CUDA; the eager
+    reference math elsewhere. The eager path is memory-bound on dozens of
+    unfused fp32 elementwise kernel passes per chunk; fusing them is a >3x
+    end-to-end step win. Same math, same dtypes — kill switch:
+    SPRIG_COMPILE=0."""
+    global _ELL_COMPILED
+    if device.type != "cuda" or os.environ.get("SPRIG_COMPILE", "1") == "0":
+        return _ell_chunk_math
+    if _ELL_COMPILED is None:
+        try:
+            # One graph per (group h/w, chunk rows, batch) shape; a process
+            # that probes several batch sizes exceeds dynamo's default limit
+            # of 8 and would silently fall back to slow eager for the rest.
+            import torch._dynamo.config as _dcfg
+            if getattr(_dcfg, "recompile_limit", 0) < 64:
+                _dcfg.recompile_limit = 64
+            if getattr(_dcfg, "cache_size_limit", 0) < 64:
+                _dcfg.cache_size_limit = 64
+            _ELL_COMPILED = torch.compile(_ell_chunk_math, dynamic=False)
+        except Exception:
+            _ELL_COMPILED = _ell_chunk_math
+    return _ELL_COMPILED
+
+
+def film_scale_shift(phi: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
+    """phi [..., 8] -> (scale [..., 3], shift [..., 3]) per RGB mean group."""
+    return 1.0 + phi[..., 0:3], phi[..., 3:6]
+
+
+def phi_at_leaf_centers(
+    Phi: torch.Tensor, rects: torch.Tensor, canvas_px: int
+) -> torch.Tensor:
+    """Bilinearly sample Phi [B, 8, 16, 16] at leaf-center canvas positions.
+
+    rects int [n, 4] (x0, y0, x1, y1 in px) -> phi [B, n, 8].
+    """
+    B = Phi.shape[0]
+    cx = (rects[:, 0] + rects[:, 2]).float() / (2.0 * canvas_px) * 2.0 - 1.0
+    cy = (rects[:, 1] + rects[:, 3]).float() / (2.0 * canvas_px) * 2.0 - 1.0
+    grid = torch.stack([cx, cy], dim=-1).to(Phi.device)          # [n, 2] (x, y)
+    grid = grid.view(1, -1, 1, 2).expand(B, -1, -1, -1)
+    out = F.grid_sample(Phi.float(), grid, mode="bilinear", align_corners=False)
+    return out.squeeze(-1).permute(0, 2, 1)                      # [B, n, 8]
+
+
+class TexelAtlas(nn.Module):
+    """cfg is duck-typed: needs T_v, d, caption_dim (optional: atlas_heads,
+    leaf_chunk)."""
+
+    def __init__(self, cfg) -> None:
+        super().__init__()
+        self.cfg = cfg
+        T_v, d = cfg.T_v, cfg.d
+        heads = getattr(cfg, "atlas_heads", 4)
+        self.leaf_chunk = int(getattr(cfg, "leaf_chunk", 16))
+
+        self.E_T = nn.Parameter(torch.randn(T_v, d) * 0.02)
+        self.q_proj = nn.Linear(d, ATLAS_WIDTH)
+        self.cap_proj = nn.Linear(cfg.caption_dim, ATLAS_WIDTH)
+        self.blocks = nn.ModuleList(
+            [CrossAttnBlock(ATLAS_WIDTH, heads) for _ in range(2)]
+        )
+        self.ln_out = nn.LayerNorm(ATLAS_WIDTH)
+        self.seed = nn.Linear(ATLAS_WIDTH, ATLAS_WIDTH * 4 * 4)
+        self.conv1 = nn.Conv2d(ATLAS_WIDTH, 128, kernel_size=3, padding=1)
+        self.conv2 = nn.Conv2d(128, dl.N_CH, kernel_size=3, padding=1)
+        # Resurrection-writable per-texel bias grid.
+        self.bias_grid = nn.Parameter(torch.zeros(T_v, dl.N_CH, ATLAS_RES, ATLAS_RES))
+
+    def render(self, emb: torch.Tensor, emb_len: torch.Tensor) -> torch.Tensor:
+        """emb [B, L, 768], emb_len [B] -> atlas [B, T_v, 40, 16, 16]."""
+        B, L, _ = emb.shape
+        T_v = self.E_T.shape[0]
+        kv = self.cap_proj(emb.float())
+        kpm = caption_padding_mask(emb_len, L, emb.device)
+
+        h = self.q_proj(self.E_T).unsqueeze(0).expand(B, -1, -1)  # [B, T_v, 256]
+        for blk in self.blocks:
+            h = blk(h, kv, kpm)
+        h = self.ln_out(h)
+
+        z = self.seed(h).reshape(B * T_v, ATLAS_WIDTH, 4, 4)
+        z = F.gelu(self.conv1(z))
+        z = F.interpolate(z, scale_factor=2, mode="nearest")      # [.., 128, 8, 8]
+        z = self.conv2(z)
+        z = F.interpolate(z, scale_factor=2, mode="nearest")      # [.., 40, 16, 16]
+        atlas = z.reshape(B, T_v, dl.N_CH, ATLAS_RES, ATLAS_RES)
+        return atlas + self.bias_grid.unsqueeze(0)
+
+    def score_leaves(
+        self,
+        atlas: torch.Tensor,
+        images: torch.Tensor,
+        lattice,
+        Phi: torch.Tensor,
+        chunk: Optional[int] = None,
+        pix_weight: Optional[torch.Tensor] = None,
+    ) -> torch.Tensor:
+        """Emission log-likelihoods for every (leaf region, texel) pair.
+
+        atlas [B, T_v, 40, 16, 16], images u8 [B, C, C, 3] (C = canvas),
+        Phi [B, 8, 16, 16] -> ell fp32 [B, n_leaf_regions, T_v], summed over
+        the region's pixels, in lattice.leaf_ids slot order. Vectorized per
+        leaf shape group, chunked over regions to bound memory.
+
+        pix_weight (optional) fp32 [B, C, C]: per-pixel importance weights on
+        the log-likelihood (object-pixel up-weighting during training). None
+        means all-ones (exact NLL) — used by log_marginal/eval/parsing.
+        """
+        B, T_v = atlas.shape[0], atlas.shape[1]
+        canvas = lattice.canvas_px
+        step = int(chunk) if chunk is not None else self.leaf_chunk
+
+        x = dl.u8_to_unit(images).permute(0, 3, 1, 2).contiguous()   # [B,3,C,C]
+        x_flat = x.reshape(B, 3, canvas * canvas)
+        if pix_weight is None:
+            w_flat = torch.ones(B, canvas * canvas, dtype=torch.float32,
+                                device=atlas.device)
+        else:
+            w_flat = pix_weight.to(atlas.device, torch.float32).reshape(
+                B, canvas * canvas)
+
+        n_leaf = lattice.n_leaf_regions
+        ell = torch.zeros(B, n_leaf, T_v, dtype=torch.float32, device=atlas.device)
+
+        leaf_rects = lattice.regions[lattice.leaf_ids]
+        phi_leaf = phi_at_leaf_centers(Phi, leaf_rects, canvas)      # [B,n_leaf,8]
+        scale, shift = film_scale_shift(phi_leaf)                    # [B,n_leaf,3]
+
+        use_ckpt = torch.is_grad_enabled() and (
+            atlas.requires_grad or Phi.requires_grad)
+        chunk_fn = _ell_chunk_fn(atlas.device)
+
+        groups = getattr(lattice, "shape_groups", None)
+        if groups is None:  # foreign/stub lattice: derive on the fly
+            groups = []
+            for (h, w), slots in lattice.leaf_shape_groups().items():
+                slots = slots.to(atlas.device)
+                rects = lattice.regions[lattice.leaf_ids[slots]]     # [n_g,4]
+                dy = torch.arange(h, device=atlas.device)
+                dx = torch.arange(w, device=atlas.device)
+                yy = rects[:, 1].view(-1, 1, 1).to(atlas.device) + dy.view(1, -1, 1)
+                xx = rects[:, 0].view(-1, 1, 1).to(atlas.device) + dx.view(1, 1, -1)
+                groups.append((h, w, slots, (yy * canvas + xx).reshape(-1)))
+
+        for h, w, slots, pix in groups:
+            if (h, w) == (ATLAS_RES, ATLAS_RES):
+                pooled = atlas                                       # identity pool
+            else:
+                pooled = F.adaptive_avg_pool2d(
+                    atlas.reshape(B * T_v, dl.N_CH, ATLAS_RES, ATLAS_RES), (h, w)
+                ).reshape(B, T_v, dl.N_CH, h, w)
+
+            n_g = slots.shape[0]
+            crops = x_flat[:, :, pix].reshape(B, 3, n_g, h, w).permute(0, 2, 1, 3, 4)
+            w_crops = w_flat[:, pix].reshape(B, n_g, h, w)
+
+            sc = scale[:, slots]                                     # [B,n_g,3]
+            sh = shift[:, slots]
+
+            for i0 in range(0, n_g, step):
+                i1 = min(n_g, i0 + step)
+                args = (pooled, crops[:, i0:i1].contiguous(),
+                        w_crops[:, i0:i1].contiguous(),
+                        sc[:, i0:i1].contiguous(), sh[:, i0:i1].contiguous())
+                if use_ckpt:
+                    tot = torch.utils.checkpoint.checkpoint(
+                        chunk_fn, *args, use_reentrant=False)
+                else:
+                    tot = chunk_fn(*args)
+                ell[:, slots[i0:i1], :] = tot
+
+        return ell

sprig/model/dl.py

@@ -0,0 +1,161 @@
+"""Discretized-logistic mixture emissions (PixelCNN++ style) with RGB coupling.
+
+Parameter layout (the "40 channels", DESIGN.md section 4): 4 mixture components,
+each a contiguous block of 10 channels, so channel index = 10*j + t for
+component j in 0..3 and t in:
+
+    t = 0      mixture-weight logit
+    t = 1..3   raw means for (R, G, B)                 (image scaled to [-1, 1])
+    t = 4..6   log-scales for (R, G, B)                (clamped to [-7, 2])
+    t = 7..9   channel-coupling coeffs (alpha: G|R, beta: B|R, gamma: B|G),
+               squashed by tanh inside this module.
+
+Coupled means (per component): m_R = mu_R;  m_G = mu_G + alpha * x_R;
+m_B = mu_B + beta * x_R + gamma * x_G, where x_R/x_G are OBSERVED values in
+[-1, 1] (teacher-forced coupling, exactly as in PixelCNN++).
+
+All math is done in fp32 regardless of input dtype.
+"""
+from __future__ import annotations
+
+import math
+from typing import Tuple
+
+import torch
+import torch.nn.functional as F
+
+N_COMP = 4
+CH_PER_COMP = 10
+N_CH = N_COMP * CH_PER_COMP  # 40
+
+WEIGHT_IDX = [10 * j for j in range(N_COMP)]
+MEAN_IDX = [10 * j + 1 + c for j in range(N_COMP) for c in range(3)]
+LOGSCALE_IDX = [10 * j + 4 + c for j in range(N_COMP) for c in range(3)]
+COEFF_IDX = [10 * j + 7 + c for j in range(N_COMP) for c in range(3)]
+
+LOGSCALE_MIN = -7.0
+LOGSCALE_MAX = 2.0
+_HALF_BIN = 1.0 / 255.0  # half bin width in [-1, 1] scale (256 levels)
+
+
+def unpack_dl_params(
+    params: torch.Tensor,
+) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor, torch.Tensor]:
+    """Split raw params [..., 40, H, W] into mixture pieces.
+
+    Returns (logit_w [..., 4, H, W], means [..., 4, 3, H, W],
+    log_scales [..., 4, 3, H, W] clamped, coeffs [..., 4, 3, H, W] tanh-ed).
+    """
+    if params.shape[-3] != N_CH:
+        raise ValueError("expected %d param channels, got %d" % (N_CH, params.shape[-3]))
+    p = params.float()
+    lead = p.shape[:-3]
+    hw = p.shape[-2:]
+    p = p.reshape(lead + (N_COMP, CH_PER_COMP) + hw)
+    logit_w = p[..., 0, :, :]
+    means = p[..., 1:4, :, :]
+    log_scales = p[..., 4:7, :, :].clamp(LOGSCALE_MIN, LOGSCALE_MAX)
+    coeffs = torch.tanh(p[..., 7:10, :, :])
+    return logit_w, means, log_scales, coeffs
+
+
+def couple_means(
+    means: torch.Tensor, coeffs: torch.Tensor, x: torch.Tensor
+) -> torch.Tensor:
+    """Coupled per-component means given observed pixels.
+
+    means/coeffs: [..., 4, 3, H, W]; x: broadcastable to [..., 3, H, W],
+    values in [-1, 1]. Returns coupled means [..., 4, 3, H, W].
+    """
+    xb = x.float().unsqueeze(-4)  # [..., 1, 3, H, W]
+    x_r = xb[..., 0, :, :]
+    x_g = xb[..., 1, :, :]
+    m_r = means[..., 0, :, :]
+    m_g = means[..., 1, :, :] + coeffs[..., 0, :, :] * x_r
+    m_b = means[..., 2, :, :] + coeffs[..., 1, :, :] * x_r + coeffs[..., 2, :, :] * x_g
+    m_r, m_g, m_b = torch.broadcast_tensors(m_r, m_g, m_b)
+    return torch.stack([m_r, m_g, m_b], dim=-3)
+
+
+def dl_channel_logprobs(
+    means: torch.Tensor,
+    log_scales: torch.Tensor,
+    coeffs: torch.Tensor,
+    x: torch.Tensor,
+) -> torch.Tensor:
+    """Per-component, per-channel discretized-logistic log P(bin of x).
+
+    means/log_scales/coeffs: [..., 4, 3, H, W] (log_scales already clamped,
+    coeffs already tanh-ed — i.e. outputs of unpack_dl_params). x broadcastable
+    to [..., 3, H, W] in [-1, 1] on the 256-level grid. Returns [..., 4, 3, H, W].
+    Edge bins (0 and 255) integrate the full tails, so per-channel bin
+    probabilities sum to 1.
+    """
+    xb = x.float().unsqueeze(-4)
+    m = couple_means(means, coeffs, x)
+    centered = xb - m
+    inv_std = torch.exp(-log_scales)
+    plus_in = inv_std * (centered + _HALF_BIN)
+    min_in = inv_std * (centered - _HALF_BIN)
+
+    log_cdf_plus = plus_in - F.softplus(plus_in)          # log sigmoid(plus_in)
+    log_one_minus_cdf_min = -F.softplus(min_in)           # log(1 - sigmoid(min_in))
+    cdf_delta = torch.sigmoid(plus_in) - torch.sigmoid(min_in)
+    mid_in = inv_std * centered
+    log_pdf_mid = mid_in - log_scales - 2.0 * F.softplus(mid_in)
+
+    log_prob_mid = torch.where(
+        cdf_delta > 1e-5,
+        torch.log(cdf_delta.clamp(min=1e-12)),
+        log_pdf_mid - math.log(127.5),
+    )
+    logp = torch.where(
+        xb < -0.999,
+        log_cdf_plus,
+        torch.where(xb > 0.999, log_one_minus_cdf_min, log_prob_mid),
+    )
+    return logp
+
+
+def dl_logprob(params: torch.Tensor, x: torch.Tensor) -> torch.Tensor:
+    """Full mixture log-prob per pixel.
+
+    params [..., 40, H, W] raw; x broadcastable to [..., 3, H, W] in [-1, 1].
+    Returns fp32 [..., H, W]: log p(x_pixel) = logsumexp_j (log w_j +
+    sum_channels log P_j(channel)).
+    """
+    logit_w, means, log_scales, coeffs = unpack_dl_params(params)
+    ch = dl_channel_logprobs(means, log_scales, coeffs, x)      # [..., 4, 3, H, W]
+    per_comp = ch.sum(dim=-3)                                    # [..., 4, H, W]
+    log_w = F.log_softmax(logit_w, dim=-3)
+    return torch.logsumexp(log_w + per_comp, dim=-3)
+
+
+def dl_mean_pixels(params: torch.Tensor) -> torch.Tensor:
+    """Deterministic 'DL-mean' rendering: per pixel pick the argmax-weight
+    component, then roll out the coupled means sequentially
+    (R = mu_R, G = mu_G + a R, B = mu_B + b R + c G), clamped to [-1, 1].
+
+    params [..., 40, H, W] -> pixels [..., 3, H, W] fp32 in [-1, 1].
+    """
+    logit_w, means, log_scales, coeffs = unpack_dl_params(params)
+    idx = logit_w.argmax(dim=-3)                                 # [..., H, W]
+    gather_idx = idx.unsqueeze(-3).unsqueeze(-4).expand(
+        means.shape[:-4] + (1, 3) + means.shape[-2:]
+    )
+    m = means.gather(-4, gather_idx).squeeze(-4)                 # [..., 3, H, W]
+    cf = coeffs.gather(-4, gather_idx).squeeze(-4)
+    r = m[..., 0, :, :].clamp(-1.0, 1.0)
+    g = (m[..., 1, :, :] + cf[..., 0, :, :] * r).clamp(-1.0, 1.0)
+    b = (m[..., 2, :, :] + cf[..., 1, :, :] * r + cf[..., 2, :, :] * g).clamp(-1.0, 1.0)
+    return torch.stack([r, g, b], dim=-3)
+
+
+def u8_to_unit(x_u8: torch.Tensor) -> torch.Tensor:
+    """uint8 [0,255] -> fp32 [-1, 1] on the 256-level grid."""
+    return x_u8.float() / 127.5 - 1.0
+
+
+def unit_to_u8(x: torch.Tensor) -> torch.Tensor:
+    """fp32 [-1, 1] -> uint8, rounding to the nearest of the 256 levels."""
+    return ((x.clamp(-1.0, 1.0) + 1.0) * 127.5).round().clamp(0, 255).to(torch.uint8)

sprig/model/gmt.py

@@ -0,0 +1,139 @@
+"""Grammar Modulation Transformer (DESIGN.md section 4, gmt.py).
+
+Symbol embeddings are queries; caption tokens are keys/values. There is no
+symbol-symbol self-attention. Outputs the caption-conditioned quantities the
+grammar needs: H, U (p(k|A,c) logits), cut-type logits, illumination field Phi,
+plus the static tables V, W, P_T and the factorized termination head.
+"""
+from __future__ import annotations
+
+from typing import NamedTuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+N_CUT_TYPES = 14
+
+
+class CrossAttnBlock(nn.Module):
+    """Pre-LN {cross-attention(queries -> memory), FFN width*4} block."""
+
+    def __init__(self, width: int, n_heads: int) -> None:
+        super().__init__()
+        self.ln_q = nn.LayerNorm(width)
+        self.ln_kv = nn.LayerNorm(width)
+        self.attn = nn.MultiheadAttention(width, n_heads, batch_first=True)
+        self.ln_f = nn.LayerNorm(width)
+        self.ffn = nn.Sequential(
+            nn.Linear(width, 4 * width), nn.GELU(), nn.Linear(4 * width, width)
+        )
+
+    def forward(
+        self,
+        q: torch.Tensor,
+        kv: torch.Tensor,
+        key_padding_mask: torch.Tensor,
+    ) -> torch.Tensor:
+        kvn = self.ln_kv(kv)
+        a, _ = self.attn(
+            self.ln_q(q), kvn, kvn,
+            key_padding_mask=key_padding_mask, need_weights=False,
+        )
+        h = q + a
+        return h + self.ffn(self.ln_f(h))
+
+
+def caption_padding_mask(emb_len: torch.Tensor, L: int, device: torch.device) -> torch.Tensor:
+    """bool [B, L], True = padding (ignored by attention). Position 0 is always
+    kept valid as a guard against fully-masked rows (NaN attention)."""
+    idx = torch.arange(L, device=device)
+    mask = idx.unsqueeze(0) >= emb_len.to(device).long().unsqueeze(1)
+    mask[:, 0] = False
+    return mask
+
+
+class GMTOut(NamedTuple):
+    H: torch.Tensor            # [B, S, d]
+    U: torch.Tensor            # [B, S, R]   p(k|A,c) logits
+    cut_logits: torch.Tensor   # [B, R, 14]  cut-type logits per component
+    Phi: torch.Tensor          # [B, 8, 16, 16] illumination field
+
+
+class GrammarModulationTransformer(nn.Module):
+    """cfg is duck-typed: needs S, R, T_v, d, n_heads, d_t, caption_dim, n_geom."""
+
+    def __init__(self, cfg) -> None:
+        super().__init__()
+        self.cfg = cfg
+        S, R, T_v, d = cfg.S, cfg.R, cfg.T_v, cfg.d
+        if d % cfg.n_heads != 0:
+            raise ValueError("d must be divisible by n_heads")
+
+        self.E_N = nn.Parameter(torch.randn(S, d) * 0.02)
+        self.cap_proj = nn.Linear(cfg.caption_dim, d)
+        self.blocks = nn.ModuleList(
+            [CrossAttnBlock(d, cfg.n_heads) for _ in range(4)]
+        )
+        self.ln_out = nn.LayerNorm(d)
+
+        # Heads.
+        self.W_u = nn.Linear(d, R)
+        self.mlp_h = nn.Sequential(nn.Linear(d, cfg.d_t), nn.GELU(), nn.Linear(cfg.d_t, cfg.d_t))
+        self.mlp_g = nn.Sequential(nn.Linear(cfg.n_geom, cfg.d_t), nn.GELU(), nn.Linear(cfg.d_t, cfg.d_t))
+        self.term_bias = nn.Parameter(torch.zeros(S))
+
+        # Cut-type head: component embeddings cross-attend once to the caption.
+        self.e_k = nn.Parameter(torch.randn(R, d) * 0.02)
+        self.cut_block = CrossAttnBlock(d, cfg.n_heads)
+        self.cut_ln = nn.LayerNorm(d)
+        self.cut_out = nn.Linear(d, N_CUT_TYPES)
+
+        # Static grammar tables (caption-independent).
+        self.P_T = nn.Parameter(torch.randn(R, T_v) * 0.02)
+        self.V = nn.Parameter(torch.randn(R, S) * 0.02)
+        self.W = nn.Parameter(torch.randn(R, S) * 0.02)
+
+        # Illumination field: masked-mean-pooled caption -> MLP -> 2x deconvs.
+        self.phi_mlp = nn.Sequential(nn.Linear(d, 128), nn.GELU(), nn.Linear(128, 32 * 4 * 4))
+        self.phi_deconv1 = nn.ConvTranspose2d(32, 16, kernel_size=4, stride=2, padding=1)
+        self.phi_deconv2 = nn.ConvTranspose2d(16, 8, kernel_size=4, stride=2, padding=1)
+        # Zero-init the final deconv so Phi == 0 at init (FiLM starts at identity).
+        nn.init.zeros_(self.phi_deconv2.weight)
+        nn.init.zeros_(self.phi_deconv2.bias)
+
+    def forward(self, emb: torch.Tensor, emb_len: torch.Tensor) -> GMTOut:
+        """emb [B, L, 768] (any float dtype), emb_len [B] int -> GMTOut."""
+        B, L, _ = emb.shape
+        kv = self.cap_proj(emb.float())                       # [B, L, d]
+        kpm = caption_padding_mask(emb_len, L, emb.device)
+
+        h = self.E_N.unsqueeze(0).expand(B, -1, -1)
+        for blk in self.blocks:
+            h = blk(h, kv, kpm)
+        H = self.ln_out(h)                                    # [B, S, d]
+
+        U = self.W_u(H)                                       # [B, S, R]
+
+        q_k = self.e_k.unsqueeze(0).expand(B, -1, -1)
+        hk = self.cut_block(q_k, kv, kpm)
+        cut_logits = self.cut_out(self.cut_ln(hk))            # [B, R, 14]
+
+        keep = (~kpm).float().unsqueeze(-1)                   # [B, L, 1]
+        pooled = (kv * keep).sum(1) / keep.sum(1).clamp(min=1.0)
+        z = self.phi_mlp(pooled).reshape(B, 32, 4, 4)
+        z = F.gelu(self.phi_deconv1(z))
+        Phi = self.phi_deconv2(z)                             # [B, 8, 16, 16]
+
+        return GMTOut(H=H, U=U, cut_logits=cut_logits, Phi=Phi)
+
+    def termination_logits(self, H: torch.Tensor, phi_geom: torch.Tensor) -> torch.Tensor:
+        """Factorized termination head:
+        term_logit[b, r, A] = MLP_h(H)[b, A, :] . MLP_g(phi_geom)[r, :] + bias_A.
+
+        H [B, S, d], phi_geom [N_reg, n_geom] -> logits [B, N_reg, S] fp32.
+        """
+        hs = self.mlp_h(H)                                    # [B, S, d_t]
+        gs = self.mlp_g(phi_geom.float())                     # [N_reg, d_t]
+        term = torch.einsum("bsd,nd->bns", hs.float(), gs)
+        return term + self.term_bias.view(1, 1, -1)

sprig/model/sprig.py

@@ -0,0 +1,593 @@
+"""SPRIGModel — wires GMT + TexelAtlas + inside DP (contracts C2/C3/C4).
+
+DP access: all calls to sprig/dp/inside.py go through the small adapters at the
+bottom of this file (`_dp`, `_call_dp`), keyed by the DESIGN.md section 5
+argument names (ell_leaf, term_logits, cut_logits, U_logmix, logV, logW,
+lattice, temper_kappa; plus log_PT for the texel prior). Tests may override the
+DP module via `_DP_MODULE`.
+
+Conventions:
+- Axiom (root) nonterminal is symbol 0.
+- Rule-logit temperature `tau` (buffer) divides all rule/termination logits
+  (U, cut-type, V, W, P_T, termination) per DESIGN.md section 2.
+- Emission tempering `eta` (buffer): kappa(r) = max(1, area_px(r)^eta);
+  `log_marginal(..., report_mode=True)` forces eta = 0 (C2).
+"""
+from __future__ import annotations
+
+import inspect
+import math
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from sprig.dp.lattice import Lattice, get_lattice
+from sprig.model import dl
+from sprig.model.atlas import ATLAS_RES, TexelAtlas, film_scale_shift, phi_at_leaf_centers
+from sprig.model.gmt import GrammarModulationTransformer
+
+_MATERIAL_SEED_OFFSET = 0x9E3779B9
+
+
+@dataclass
+class SPRIGConfig:
+    S: int = 1024
+    R: int = 64
+    T_v: int = 256
+    d: int = 384
+    canvas: int = 64
+    grid: int = 8
+    leaf_max: int = 16
+    n_heads: int = 6
+    d_t: int = 64
+    caption_dim: int = 768
+    n_geom: int = 64
+    atlas_heads: int = 4
+    leaf_chunk: int = 16
+    axiom: int = 0
+    texel_hinge_weight: float = 1.0
+    symbol_hinge_weight: float = 0.5
+    # Object-pixel importance weight on the emission log-likelihood during
+    # training (diagnosis fix 1: objects are <10% of pixels, so unweighted NLL
+    # lets background fidelity dominate and no object texels ever form).
+    # Requires the dataloader to provide batch["objmask"] u8 [B,C,C].
+    # 1.0 = exact NLL (eval/log_marginal always use exact NLL regardless).
+    emission_obj_weight: float = 1.0
+    # Second-order grad through the DP so the under-use hinges actually train
+    # (usage comes from autograd.grad of logZ; without create_graph the hinge
+    # would be constant w.r.t. parameters).
+    hinge_create_graph: bool = True
+    resurrect_threshold_frac: float = 0.1  # of uniform usage 1/T_v
+
+    def __post_init__(self) -> None:
+        if self.d % self.n_heads != 0:
+            raise ValueError("d must be divisible by n_heads")
+
+
+@dataclass
+class ParseNode:
+    rect: Tuple[int, int, int, int]
+    axis: Optional[int]          # None for leaves; 0 = vertical cut, 1 = horizontal
+    cut_px: Optional[int]
+    symbol: int
+    texel: Optional[int]         # None for internal nodes
+    children: List["ParseNode"] = field(default_factory=list)
+
+
+class SPRIGModel(nn.Module):
+    def __init__(self, cfg: SPRIGConfig) -> None:
+        super().__init__()
+        self.cfg = cfg
+        self.gmt = GrammarModulationTransformer(cfg)
+        self.atlas = TexelAtlas(cfg)
+        self.lattice: Lattice = get_lattice(cfg.canvas, cfg.grid, cfg.leaf_max)
+        self.register_buffer("tau", torch.tensor(1.0))
+        self.register_buffer("eta", torch.tensor(0.0))
+
+    # ------------------------------------------------------------------ utils
+
+    def _lat(self, device: torch.device) -> Lattice:
+        self.lattice.to(device)
+        return self.lattice
+
+    def _kappa(self, eta: float, device: torch.device) -> torch.Tensor:
+        """kappa(r) = max(1, area_px(r)^eta) over leaf slots -> fp32 [n_leaf]."""
+        lat = self._lat(device)
+        area = lat.area_px[lat.leaf_ids].float()
+        if eta <= 0.0:
+            return torch.ones_like(area)
+        return torch.clamp(area ** eta, min=1.0)
+
+    def _conditionals(
+        self,
+        emb: torch.Tensor,
+        emb_len: torch.Tensor,
+        images: Optional[torch.Tensor] = None,
+        pix_weight: Optional[torch.Tensor] = None,
+    ) -> Dict[str, torch.Tensor]:
+        """All caption-conditioned quantities, temperature already applied."""
+        lat = self._lat(emb.device)
+        out = self.gmt(emb, emb_len)
+        tau = self.tau.clamp(min=1e-3)
+        cond: Dict[str, torch.Tensor] = {}
+        cond["H"] = out.H
+        cond["Phi"] = out.Phi
+        cond["U_logmix"] = F.log_softmax(out.U / tau, dim=-1)          # [B,S,R]
+        cond["log_PT"] = F.log_softmax(self.gmt.P_T / tau, dim=-1)     # [R,T_v]
+        cond["logV"] = F.log_softmax(self.gmt.V / tau, dim=-1)         # [R,S]
+        cond["logW"] = F.log_softmax(self.gmt.W / tau, dim=-1)         # [R,S]
+        cond["term_logits"] = self.gmt.termination_logits(out.H, lat.phi_geom) / tau
+        cond["cut_logits"] = out.cut_logits / tau                      # [B,R,14]
+        cond["atlas"] = self.atlas.render(emb, emb_len)
+        if images is not None:
+            cond["ell"] = self.atlas.score_leaves(
+                cond["atlas"], images, lat, out.Phi, pix_weight=pix_weight)
+        return cond
+
+    def _texel_prior_log(self, cond: Dict[str, torch.Tensor]) -> torch.Tensor:
+        """log p(T|A,c) = logsumexp_k(log p(k|A,c) + log p(T|k)) -> [B,S,T_v]."""
+        return torch.logsumexp(
+            cond["U_logmix"].unsqueeze(-1) + cond["log_PT"].unsqueeze(0).unsqueeze(0),
+            dim=2,
+        )
+
+    def _dp_kwargs(
+        self, cond: Dict[str, torch.Tensor], kappa: torch.Tensor
+    ) -> Dict[str, Any]:
+        return dict(
+            ell_leaf=cond["ell"],
+            term_logits=cond["term_logits"],
+            cut_logits=cond["cut_logits"],
+            U_logmix=cond["U_logmix"],
+            logV=cond["logV"],
+            logW=cond["logW"],
+            lattice=self.lattice,
+            temper_kappa=kappa,
+            log_PT=cond["log_PT"],
+        )
+
+    def _inside(self, cond: Dict[str, torch.Tensor], kappa: torch.Tensor) -> torch.Tensor:
+        mod = _dp()
+        fn = getattr(mod, "inside_logZ", None) or getattr(mod, "inside")
+        res = _call_dp(fn, self._dp_kwargs(cond, kappa))
+        if isinstance(res, (tuple, list)):
+            return res[-1]  # (beta, logZ) per DESIGN section 5
+        return res
+
+    # -------------------------------------------------------------- contracts
+
+    def log_marginal(
+        self,
+        image: torch.Tensor,
+        emb: torch.Tensor,
+        emb_len: torch.Tensor,
+        report_mode: bool = False,
+    ) -> torch.Tensor:
+        """C2: exact log p(x|c) [B]; report_mode forces eta = 0."""
+        eta = 0.0 if report_mode else float(self.eta)
+        cond = self._conditionals(emb, emb_len, images=image)
+        kappa = self._kappa(eta, emb.device)
+        return self._inside(cond, kappa)
+
+    def loss(self, batch: Dict[str, torch.Tensor]) -> Tuple[torch.Tensor, Dict[str, float]]:
+        """DESIGN section 6: tempered NLL (nats/subpixel) + under-use hinges.
+
+        Usage vectors are posterior expected counts obtained from
+        torch.autograd.grad(logZ.sum(), [ell, U_logmix], retain_graph=True):
+        grad wrt ell (x kappa) = expected (leaf, texel) counts; grad wrt
+        U_logmix summed over k = expected node counts per symbol (the texel
+        prior is mixed from U_logmix, so termination counts are included).
+        """
+        images = batch["image"]
+        emb = batch["emb"]
+        emb_len = batch["emb_len"]
+        cfg = self.cfg
+        pix_weight = None
+        if cfg.emission_obj_weight != 1.0 and batch.get("objmask") is not None:
+            pix_weight = 1.0 + (cfg.emission_obj_weight - 1.0) * batch[
+                "objmask"].to(images.device, torch.float32)
+        cond = self._conditionals(emb, emb_len, images=images,
+                                  pix_weight=pix_weight)
+        kappa = self._kappa(float(self.eta), emb.device)
+        logZ = self._inside(cond, kappa)
+
+        n_subpix = 3 * cfg.canvas * cfg.canvas
+        nll = (-logZ / n_subpix).mean()
+
+        create = cfg.hinge_create_graph and torch.is_grad_enabled()
+        g_ell, g_u = torch.autograd.grad(
+            logZ.sum(),
+            [cond["ell"], cond["U_logmix"]],
+            retain_graph=True,
+            create_graph=create,
+        )
+        texel_counts = (g_ell * kappa.view(1, -1, 1)).sum(dim=(0, 1))   # [T_v]
+        symbol_counts = g_u.sum(dim=(0, 2))                             # [S]
+        texel_usage = texel_counts / texel_counts.sum().clamp(min=1e-12)
+        symbol_usage = symbol_counts / symbol_counts.sum().clamp(min=1e-12)
+
+        texel_hinge = F.relu(1.0 / (4.0 * cfg.T_v) - texel_usage).sum()
+        symbol_hinge = F.relu(1.0 / (4.0 * cfg.S) - symbol_usage).sum()
+        total = (
+            nll
+            + cfg.texel_hinge_weight * texel_hinge
+            + cfg.symbol_hinge_weight * symbol_hinge
+        )
+
+        with torch.no_grad():
+            tu = texel_usage.detach()
+            su = symbol_usage.detach()
+            metrics = {
+                "loss": float(total),
+                "nll": float(nll),
+                "bpd": float(nll) / math.log(2.0),
+                "logZ_mean": float(logZ.mean()),
+                "texel_hinge": float(texel_hinge),
+                "symbol_hinge": float(symbol_hinge),
+                "texel_alive_frac": float((tu >= 1.0 / (4.0 * cfg.T_v)).float().mean()),
+                "symbol_eff": float(torch.exp(-(su.clamp(min=1e-12) * su.clamp(min=1e-12).log()).sum())),
+                "mean_leaves": float((g_ell.detach() * kappa.view(1, -1, 1)).sum() / images.shape[0]),
+            }
+        return total, metrics
+
+    def posterior_usage(
+        self, image: torch.Tensor, emb: torch.Tensor, emb_len: torch.Tensor
+    ) -> Dict[str, Any]:
+        """C3: expected-count diagnostics from posterior marginals.
+
+        node_entropy = occupancy-weighted entropy of the conditional
+        split posterior at each region (choices: terminate + each concrete
+        cut). mean_depth is the area-based proxy E[log2(canvas_area /
+        leaf_area)] over leaf marginals.
+        """
+        cfg = self.cfg
+        with torch.enable_grad():
+            cond = self._conditionals(emb, emb_len, images=image)
+            kappa = self._kappa(float(self.eta), emb.device)
+            marg = _call_dp(_dp().posterior_marginals, self._dp_kwargs(cond, kappa))
+
+        lat = self._lat(emb.device)
+        node = marg["node"].float()          # [B, N_reg, S]
+        term = marg["term"].float()          # [B, N_reg, S]
+        cut = marg["cut"].float()            # [B, M_total] concatenated level rows
+        texel = marg["texel"].float()        # [B, n_leaf, T_v]
+        rule = marg["rule"].float()          # [B, S, R]
+
+        symbol_counts = node.sum(dim=(0, 1))
+        texel_counts = texel.sum(dim=(0, 1))
+        symbol_usage = symbol_counts / symbol_counts.sum().clamp(min=1e-12)
+        texel_usage = texel_counts / texel_counts.sum().clamp(min=1e-12)
+
+        # Occupancy-weighted entropy of conditional split posteriors.
+        B = node.shape[0]
+        occ = node.sum(dim=-1)               # [B, N_reg]
+        term_tot = term.sum(dim=-1)          # [B, N_reg]
+        parent_of_row = torch.cat([lv.parent_ids for lv in lat.levels]).to(emb.device)
+        eps = 1e-12
+        occ_safe = occ.clamp(min=eps)
+        q_term = (term_tot / occ_safe).clamp(min=0.0)
+        plogp_term = torch.where(
+            q_term > eps, q_term * q_term.clamp(min=eps).log(), torch.zeros_like(q_term)
+        )
+        q_cut = cut / occ_safe[:, parent_of_row]
+        plogp_cut_rows = torch.where(
+            q_cut > eps, q_cut * q_cut.clamp(min=eps).log(), torch.zeros_like(q_cut)
+        )
+        plogp_cut = torch.zeros_like(occ).index_add_(
+            1, parent_of_row, plogp_cut_rows
+        )
+        ent_r = -(plogp_term + plogp_cut)    # [B, N_reg]
+        w = occ.clamp(min=0.0)
+        node_entropy = float((w * ent_r).sum() / w.sum().clamp(min=eps))
+
+        # Magnitude monitors.
+        ell = cond["ell"].detach().float()
+        emit_mag = float(
+            ((texel * ell).sum() / texel.sum().clamp(min=eps)).abs()
+        )
+        rule_mag = float(
+            ((rule * cond["U_logmix"].detach()).sum() / rule.sum().clamp(min=eps)).abs()
+        )
+
+        leaf_area = lat.area_px[lat.leaf_ids].float()
+        depth_proxy = torch.log2(float(cfg.canvas * cfg.canvas) / leaf_area)
+        leaf_marg = texel.sum(dim=-1)        # [B, n_leaf]
+        mean_depth = float(
+            (leaf_marg * depth_proxy.unsqueeze(0)).sum() / leaf_marg.sum().clamp(min=eps)
+        )
+        mean_leaves = float(leaf_marg.sum() / B)
+
+        return {
+            "symbol_usage": symbol_usage.detach(),
+            "texel_usage": texel_usage.detach(),
+            "node_entropy": node_entropy,
+            "emit_mag": emit_mag,
+            "rule_mag": rule_mag,
+            "mean_depth": mean_depth,
+            "mean_leaves": mean_leaves,
+        }
+
+    def map_parse(
+        self, image: torch.Tensor, emb: torch.Tensor, emb_len: torch.Tensor
+    ) -> List[ParseNode]:
+        """C3: Viterbi (max-semiring) parse per batch element."""
+        with torch.no_grad():
+            cond = self._conditionals(emb, emb_len, images=image)
+            kappa = self._kappa(float(self.eta), emb.device)
+            res = _call_dp(_dp().viterbi, self._dp_kwargs(cond, kappa))
+        _score, trees = res
+        return [self._to_parse_node(t) for t in trees]
+
+    def _to_parse_node(self, node: Tuple) -> ParseNode:
+        """Convert the DP viterbi tree tuple (region_id, symbol, texel, axis,
+        cut_px, children) into a ParseNode."""
+        rid, sym, texel, axis, cut_px, children = node
+        rect = tuple(int(v) for v in self.lattice.regions[rid].tolist())
+        return ParseNode(
+            rect=rect,
+            axis=None if axis is None else int(axis),
+            cut_px=None if cut_px is None else int(cut_px),
+            symbol=int(sym),
+            texel=None if texel is None else int(texel),
+            children=[self._to_parse_node(ch) for ch in children],
+        )
+
+    # --------------------------------------------------------------- sampling
+
+    @torch.no_grad()
+    def sample(
+        self,
+        emb: torch.Tensor,
+        emb_len: torch.Tensor,
+        seed_struct: int,
+        seed_material: int,
+        n: int,
+    ) -> Tuple[torch.Tensor, List[ParseNode]]:
+        """C4: ancestral sampling, breadth-parallel over the frontier.
+
+        Two RNG streams: structural draws (termination, k, cut, B, C) from
+        seed_struct; material draws (texel choice) from seed_material. Pixels
+        are rendered as DL means (no pixel noise). Returns
+        (images u8 [n, canvas, canvas, 3] on CPU, list of n ParseNode roots).
+        """
+        images, trees, _scores = self._sample_scored(emb, emb_len, seed_struct, seed_material, n)
+        return images, trees
+
+    @torch.no_grad()
+    def sample_bestof(
+        self, emb: torch.Tensor, emb_len: torch.Tensor, K: int, seed: int
+    ) -> Tuple[torch.Tensor, List[ParseNode]]:
+        """C4: sample K derivations, rerank by joint log p(tree, rendered
+        image | c), return the best as (images u8 [1, canvas, canvas, 3], [tree])."""
+        images, trees, scores = self._sample_scored(
+            emb, emb_len, int(seed), int(seed) + _MATERIAL_SEED_OFFSET, K
+        )
+        best = int(torch.tensor(scores).argmax())
+        return images[best : best + 1], [trees[best]]
+
+    @torch.no_grad()
+    def _sample_scored(
+        self,
+        emb: torch.Tensor,
+        emb_len: torch.Tensor,
+        seed_struct: int,
+        seed_material: int,
+        n: int,
+    ) -> Tuple[torch.Tensor, List[ParseNode], List[float]]:
+        cfg = self.cfg
+        device = next(self.parameters()).device
+        if emb.dim() == 2:
+            emb = emb.unsqueeze(0)
+        if not torch.is_tensor(emb_len):
+            emb_len = torch.tensor([emb_len])
+        emb_len = emb_len.reshape(-1)[:1]
+        emb = emb[:1].to(device)
+
+        cond = self._conditionals(emb, emb_len)
+        lat = self._lat(device)
+
+        # Single-caption conditionals on CPU for generator-driven draws.
+        u_logmix = cond["U_logmix"][0].float().cpu()             # [S,R]
+        term_logits = cond["term_logits"][0].float().cpu()       # [N,S]
+        cut_logits = cond["cut_logits"][0].float().cpu()         # [R,14]
+        logv = cond["logV"].float().cpu()                        # [R,S]
+        logw = cond["logW"].float().cpu()                        # [R,S]
+        texel_prior = self._texel_prior_log(cond)[0].float().cpu()  # [S,T_v]
+        atlas0 = cond["atlas"][0].float().cpu()                  # [T_v,40,16,16]
+        phi0 = cond["Phi"][:1].float().cpu()                     # [1,8,16,16]
+
+        pooled_cache: Dict[Tuple[int, int], torch.Tensor] = {}
+        for (h, w), _slots in lat.leaf_shape_groups().items():
+            pooled_cache[(h, w)] = F.adaptive_avg_pool2d(atlas0, (h, w))
+
+        gen_s = torch.Generator().manual_seed(int(seed_struct) & 0x7FFFFFFFFFFFFFFF)
+        gen_m = torch.Generator().manual_seed(int(seed_material) & 0x7FFFFFFFFFFFFFFF)
+
+        regions_cpu = lat.regions.cpu()
+        leaf_mask = lat.leaf_mask.cpu()
+        must_term = lat.must_terminate.cpu()
+        type_present = lat.type_present.cpu()
+
+        images_out: List[torch.Tensor] = []
+        trees: List[ParseNode] = []
+        scores: List[float] = []
+        for _ in range(n):
+            root = ParseNode(
+                rect=tuple(int(v) for v in regions_cpu[lat.root_id].tolist()),
+                axis=None, cut_px=None, symbol=cfg.axiom, texel=None,
+            )
+            frontier: List[Tuple[int, int, ParseNode]] = [(lat.root_id, cfg.axiom, root)]
+            leaves: List[Tuple[Tuple[int, int, int, int], int]] = []
+            logp = 0.0
+            while frontier:
+                nxt: List[Tuple[int, int, ParseNode]] = []
+                for rid, sym, pn in frontier:
+                    can_term = bool(leaf_mask[rid])
+                    forced_term = bool(must_term[rid])
+                    if forced_term:
+                        terminate = True
+                    elif not can_term:
+                        terminate = False
+                    else:
+                        p_term = torch.sigmoid(term_logits[rid, sym])
+                        u = torch.rand((), generator=gen_s)
+                        terminate = bool(u < p_term)
+                        logp += float(torch.log(p_term if terminate else 1.0 - p_term))
+                    if terminate:
+                        probs = torch.exp(texel_prior[sym])
+                        t = int(torch.multinomial(probs, 1, generator=gen_m))
+                        logp += float(texel_prior[sym, t])
+                        pn.texel = t
+                        leaves.append((pn.rect, t))
+                    else:
+                        k = int(torch.multinomial(torch.exp(u_logmix[sym]), 1, generator=gen_s))
+                        logp += float(u_logmix[sym, k])
+                        # Concrete cut: masked cut-type softmax, mass split
+                        # uniformly across same-type concrete cuts.
+                        cuts = lat.cuts_of_region[rid]
+                        tl = cut_logits[k].masked_fill(~type_present[rid], float("-inf"))
+                        tls = F.log_softmax(tl, dim=-1)
+                        row_logp = torch.tensor(
+                            [float(tls[c[4]]) - c[5] for c in cuts]
+                        )
+                        ci = int(torch.multinomial(torch.exp(row_logp), 1, generator=gen_s))
+                        logp += float(row_logp[ci])
+                        axis, px, lo, hi, _t, _lc = cuts[ci]
+                        b_sym = int(torch.multinomial(torch.exp(logv[k]), 1, generator=gen_s))
+                        c_sym = int(torch.multinomial(torch.exp(logw[k]), 1, generator=gen_s))
+                        logp += float(logv[k, b_sym]) + float(logw[k, c_sym])
+                        pn.axis, pn.cut_px = int(axis), int(px)
+                        ch_lo = ParseNode(
+                            rect=tuple(int(v) for v in regions_cpu[lo].tolist()),
+                            axis=None, cut_px=None, symbol=b_sym, texel=None,
+                        )
+                        ch_hi = ParseNode(
+                            rect=tuple(int(v) for v in regions_cpu[hi].tolist()),
+                            axis=None, cut_px=None, symbol=c_sym, texel=None,
+                        )
+                        pn.children = [ch_lo, ch_hi]
+                        nxt.append((lo, b_sym, ch_lo))
+                        nxt.append((hi, c_sym, ch_hi))
+                frontier = nxt
+
+            img, emit_ll = self._render_leaves(leaves, pooled_cache, phi0)
+            images_out.append(img)
+            trees.append(root)
+            scores.append(logp + emit_ll)
+
+        return torch.stack(images_out, dim=0), trees, scores
+
+    def _render_leaves(
+        self,
+        leaves: List[Tuple[Tuple[int, int, int, int], int]],
+        pooled_cache: Dict[Tuple[int, int], torch.Tensor],
+        phi0: torch.Tensor,
+    ) -> Tuple[torch.Tensor, float]:
+        """Paint DL-mean pixels for each leaf; also return the total emission
+        log-likelihood of the rendered (quantized) pixels — the emission part
+        of the joint score used by sample_bestof."""
+        cfg = self.cfg
+        canvas = torch.zeros(3, cfg.canvas, cfg.canvas)
+        emit_ll = 0.0
+        for (x0, y0, x1, y1), texel in leaves:
+            h, w = y1 - y0, x1 - x0
+            p = pooled_cache[(h, w)][texel].clone()              # [40,h,w]
+            rect_t = torch.tensor([[x0, y0, x1, y1]])
+            phi = phi_at_leaf_centers(phi0, rect_t, cfg.canvas)[0, 0]  # [8]
+            scale, shift = film_scale_shift(phi)
+            for c in range(3):
+                idx = [10 * j + 1 + c for j in range(dl.N_COMP)]
+                p[idx] = p[idx] * scale[c] + shift[c]
+            pix = dl.dl_mean_pixels(p.unsqueeze(0))[0]           # [3,h,w]
+            pix_u8 = dl.unit_to_u8(pix)
+            canvas[:, y0:y1, x0:x1] = dl.u8_to_unit(pix_u8)
+            emit_ll += float(
+                dl.dl_logprob(p.unsqueeze(0), dl.u8_to_unit(pix_u8).unsqueeze(0)).sum()
+            )
+        img_u8 = dl.unit_to_u8(canvas).permute(1, 2, 0).contiguous()  # [C,C,3] u8
+        return img_u8, emit_ll
+
+    # ----------------------------------------------------------- resurrection
+
+    @torch.no_grad()
+    def resurrect_texels(
+        self,
+        usage: torch.Tensor,
+        images: torch.Tensor,
+        generator: Optional[torch.Generator] = None,
+        threshold: Optional[float] = None,
+        noise_std: float = 0.01,
+        obj_mask: Optional[torch.Tensor] = None,
+    ) -> int:
+        """F3.3: overwrite the bias grid of under-used texels with a
+        training-image 16x16 crop converted to DL-mean params (small noise on
+        the other channels) and perturb the E_T row. Returns #resurrected.
+
+        usage: [T_v] normalized posterior usage; images: u8 [B, canvas, canvas, 3].
+        obj_mask (optional) [B, canvas, canvas] bool/u8: when given, crops are
+        centered on object pixels (diagnosis fix 3 — random crops are ~90%
+        background, so resurrection used to reseed dead texels with yet more
+        background material).
+        """
+        cfg = self.cfg
+        thr = threshold if threshold is not None else cfg.resurrect_threshold_frac / cfg.T_v
+        dead = torch.nonzero(usage.cpu() < thr, as_tuple=False).reshape(-1)
+        if dead.numel() == 0:
+            return 0
+        gen = generator
+        B = images.shape[0]
+        hi_y = cfg.canvas - ATLAS_RES
+        dev = self.atlas.bias_grid.device
+        obj_pix: List[Tuple[int, torch.Tensor]] = []
+        if obj_mask is not None:
+            m = obj_mask.detach().cpu()
+            for b in range(B):
+                nz = torch.nonzero(m[b] > 0, as_tuple=False)   # [k, 2] (y, x)
+                if nz.numel():
+                    obj_pix.append((b, nz))
+        for t in dead.tolist():
+            if obj_pix:
+                b, nz = obj_pix[int(torch.randint(0, len(obj_pix), (1,), generator=gen))]
+                cy, cx = nz[int(torch.randint(0, nz.shape[0], (1,), generator=gen))].tolist()
+                y0 = min(max(cy - ATLAS_RES // 2, 0), hi_y)
+                x0 = min(max(cx - ATLAS_RES // 2, 0), hi_y)
+            else:
+                b = int(torch.randint(0, B, (1,), generator=gen))
+                y0 = int(torch.randint(0, hi_y + 1, (1,), generator=gen))
+                x0 = int(torch.randint(0, hi_y + 1, (1,), generator=gen))
+            crop = images[b, y0 : y0 + ATLAS_RES, x0 : x0 + ATLAS_RES].cpu()
+            crop_unit = dl.u8_to_unit(crop).permute(2, 0, 1)     # [3,16,16]
+            bias = noise_std * torch.randn(dl.N_CH, ATLAS_RES, ATLAS_RES, generator=gen)
+            for j in range(dl.N_COMP):
+                bias[10 * j + 1 : 10 * j + 4] = crop_unit
+            self.atlas.bias_grid.data[t] = bias.to(dev)
+            self.atlas.E_T.data[t] += noise_std * torch.randn(
+                cfg.d, generator=gen
+            ).to(self.atlas.E_T.device)
+        return int(dead.numel())
+
+
+# ------------------------------------------------------------------ DP access
+
+_DP_MODULE: Optional[Any] = None  # test hook: module-like override
+
+
+def _dp() -> Any:
+    if _DP_MODULE is not None:
+        return _DP_MODULE
+    from sprig.dp import inside
+    return inside
+
+
+def _call_dp(fn: Any, kwargs: Dict[str, Any]) -> Any:
+    """Call a DP function by DESIGN section 5 keyword names, dropping any
+    kwargs the target does not accept (e.g. log_PT if the DP mixes the texel
+    prior itself)."""
+    sig = inspect.signature(fn)
+    if any(p.kind == inspect.Parameter.VAR_KEYWORD for p in sig.parameters.values()):
+        return fn(**kwargs)
+    accepted = {k: v for k, v in kwargs.items() if k in sig.parameters}
+    return fn(**accepted)