[ncylich]

Three transpiler correctness fixes, each reached and validated with a runnable repro on the real capture/transpile pipeline (this branch vs origin/main). Final state: 2 files, 4 lines.

Contributions (validated)

1. JAX prefill is_causal gating (jax_semantic_rewrites.py)
A maskless self-attention is treated as causal only when the component is a known causal-prefill (decoder_prefill_chunk / decoder_media_step). Previously any maskless self-attention was forced causal, so bidirectional/encoder self-attention captured via JAX was silently causal-masked.
Validated: a maskless JAX self-attention (no component set) transpiled to output matching the causal reference and differing from true full attention by ~0.27; with the gate it matches full attention (FP16 noise).

2. JAX min/max kept in native FP32 (jax_semantic_rewrites.py)
min/max are no longer legalized through FP16 (cast→reduce→cast). FP16 saturates at ±65504, so attention-mask sentinels (-1e4, finfo.min) and large logits overflow and corrupt exact min/max selection. sum/mean still legalize to FP16 (averaging at the engine working precision). Also .strip() the component-meta string.
Validated: a min/max node on FP32 gets 0 precision-casts on this branch vs 2 (forced FP16) before; sum still gets the FP16 pair.

3. Attention-output layout guard (optimize_graph.py)
normalize_attention_layouts no longer folds an attention's output permute into the op when that attention output is itself a graph output — folding left graph.outputs referencing a value no longer produced by any node (a dangling output). Reachable for models that return a raw attention tensor (encoders / feature extractors).
Validated: with the guard the graph output is preserved; without it, the output value is folded away and dangling.

Excluded after validation

The original fuse_attention mask-input preservation was dropped — proven unreachable and, where its path is reachable, cosmetic:

• fuse_attention runs post-canonicalization (op always attention) and match_attention sets q/k/v = inputs[0:3], so a non-window masked node is always idempotent. Instrumenting the mask-preserving branch across the full test suite + repros gave 0 hits.
• The only masked-node rewrite is sliding-window, where both origin and this branch correctly drop the mask input (the window is carried by window_size); the sole difference was a leftover additive_mask=True attr on origin. Kernel test (window_size=2, no mask, additive_mask True vs absent) → identical output (0.0 diff), since with no mask tensor the flag has nothing to apply to.

A correct guard would be >3 lines for no observable gain, so it was removed rather than kept as dead code.

Verification

pytest python/tests + pytest python/cactus/convert/tests → 264 passed, 1 skipped (JAX capture tests included).

[kar-m]

@jakmro looks good i think

[ncylich]

@kar-m @jakmro heads up — I force-pushed a tightened version of this PR after reproducing every change against origin/main on the real pipeline. Net: it's now 2 files / 4 lines, and I dropped one change as unreachable.

Kept (each reproduced with numerics):

• is_causal gating — a maskless self-attention with an unknown/non-prefill component was being forced causal. Repro: transpiled output matched the causal reference and differed from true full attention by ~0.27; with the gate it matches full attention.
• min/max native FP32 — they were legalized through FP16; FP16 (±65504) overflows mask sentinels (-1e4, finfo.min) and large logits, corrupting exact selection. Now native FP32; sum/mean still FP16.
• attention-output layout guard — folding the output permute when the attention output is itself a graph output left graph.outputs dangling (value no longer produced). Now skipped.

Dropped: the fuse_attention mask-input preservation. I could not reach it: fuse_attention runs post-canonicalization so the node is always canonical attention, which makes it idempotent for any non-window masked node; the only masked rewrite is sliding-window (mask intentionally dropped, carried by window_size). Instrumenting that exact branch across the whole test suite + repros gave 0 hits. A correct defensive guard would be >3 lines, so I removed it rather than keep dead code.

pytest python/tests + pytest python/cactus/convert/tests → 264 passed, 1 skipped (JAX capture tests run).

[Copilot]

Pull request overview

This PR applies three targeted transpiler correctness fixes affecting JAX attention semantics, reduction precision handling, and attention-output layout normalization safety. The changes address silent behavioral drift (incorrect causal masking), numerical overflow risks (FP16 min/max saturation), and a graph integrity bug (dangling graph.outputs after layout folding).

Changes:

• Gate is_causal for maskless JAX self-attention to only known causal-prefill components, avoiding accidental causal-masking of bidirectional/encoder attention.
• Stop FP16-legalizing JAX min/max reductions so they remain in native FP32; also normalize the component meta string via .strip().lower().
• Prevent folding an attention output permute into the attention op when the pre-permute attention output is itself a graph output (avoids dangling outputs).

Reviewed changes

Copilot reviewed 2 out of 2 changed files in this pull request and generated no comments.

File	Description
python/cactus/transpile/optimize_graph.py	Adds a guard in normalize_attention_layouts to avoid folding away values referenced by graph.outputs.
python/cactus/transpile/jax_semantic_rewrites.py	Refines JAX attention causality inference and adjusts reduction legalization to preserve FP32 for min/max.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.