[ncylich]

Surgical transpiler correctness fixes. Each change below was reached with a runnable repro on the real capture/import pipeline; changes that no model can reach were dropped (see review thread).

• aten_ops.py — match the longest op prefix first. Confirmed normalize_target on origin mis-routes aten.addmm→add, aten.minimum→min, aten.maximum→max, aten.slice_scatter→slice, aten.select_scatter→index; a model using torch.minimum/torch.maximum/torch.addmm reaches this.
• importers.py — fail closed on scaled addmm (beta/alpha != 1) instead of silently dropping the scale. Reached: torch.addmm(bias, a, b, beta=2.0) now raises a clear capture error.
• fusion/rms_norm.py — only fuse RMSNorm when the mean reduces the last dim. Reached: a channel-wise x * rsqrt(x.pow(2).mean(1, keepdim=True) + eps) * w over (B,C,H,W) (FP16) is otherwise fused into a last-dim kernel with a wrong-sized weight (observed: weight (5,) for a channel dim of 4) → silently wrong output. The guard skips that fusion while leaving normal last-dim RMSNorm fused (no regression).
• capture_jax.py — stride-aware slice aliasing + skip a duplicate gather-squeeze reshape. Reached: x[::2] on a length-6 vector returns the full length-6 input aliased as a no-op on origin (out_shape=(6,)), vs a correct length-3 slice on this branch.

Dropped as unreachable by any model on the FP16 engine / standard import pipeline (verified empirically):

• lower.py clamp/finite-scalar changes — the only behavioral change is the non-FP16 branch, which the FP16 working precision never executes.
• import_semantics.py — the SDPA nonzero-dropout guard (dropout_p is always 0 after model.eval()), and the literal-mask passthrough (attrs["mask"] is only set for a non-tensor literal mask, but SDPA attn_mask is always a Tensor/None, so it is never produced).
• importers.py BatchNorm no-training branch — _native_batch_norm_legit_no_training only appears after run_decompositions, which in this repo runs solely on the NPU→CoreML path; the IR importer always sees the 8-arg aten.batch_norm.default, handled identically before and after.

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

[kar-m]

@ncylich not sure how to reproduce this bug, dont know if theres a model that trips any of these issues and dont wanna potentially introduce bugs regressing some model that we support

[ncylich]

@kar-m good call wanting repro + regression confidence before merging. I reproduced each of these against origin/main vs this branch with the real capture+import pipeline. Short version: 5 of the 6 are latent/defensive guards — no model we support trips them today — and none can regress a supported model (proven, not asserted). The one genuinely reachable fix is lower.py's inf-materialization.

File	What trips it	Repro (origin → with PR)	Regress a supported model?
aten_ops.py	graph with aten.minimum/maximum/slice_scatter/select_scatter/addmm	origin mis-maps minimum→min (silently drops the 2nd operand), slice_scatter→slice (TypeError crash), addmm→add (drops the matmul); PR maps each correctly	No. Verified via canonical_ir_op + an exhaustive 495-op diff: only those 5 spellings change; add/min/max/slice/select/linear/matmul/clamp are byte-identical. No supported family emits those ops.
importers.py	eval BatchNorm decomposed to _native_batch_norm_legit_no_training (origin reads args[7] → IndexError), or addmm with beta/alpha≠1 (origin silently drops the scaling)	ConvBN().eval() → run_decompositions → origin IndexError, PR imports fine; scaled torch.addmm → origin silent wrong-math, PR fails closed	No. 9-arg batch_norm.default and default addmm handled identically. Default transpile never emits the no_training variant (only run_decompositions, which feeds CoreML, not our IR).
fusion/rms_norm.py	RMSNorm-shaped subgraph whose mean is over a non-last dim (origin mis-fuses → wrong norm)	synthetic non-last-dim mean+rsqrt: origin fuses (wrong), PR skips	No. Every supported RMSNorm (gemma/qwen/lfm2/…) reduces over -1 and still fuses identically.
lower.py — the reachable one	masked_fill/where(mask, x, float('-inf')) on FP32 attention scores → origin computes 0 * -inf = NaN (mask poisoning). (Also tensor-bound clamp, which is latent.)	synthetic FP32 masked-softmax: origin → NaNs, PR → correct	No. FP16 and all finite mask values are byte-identical; only the broken non-FP16 ±inf case changes. Op-reachable (generic masked attention), not pinned to a named family — modern HF often uses finfo(dtype).min, a finite value.
import_semantics.py	SDPA node exported with dropout_p != 0 (origin silently ignores it)	synthetic SDPA dropout_p=0.1: origin drops it, PR raises	No. dropout_p=0 (every eval-mode export) is unchanged.
capture_jax.py	JAX-path lax.slice stride≠1 / duplicate gather-squeeze	only via jax_user_graph_bundle	No. JAX-only path; the default transpile is PyTorch and jax isn't installed. No supported model reaches it.

On the regression worry directly: the new code in every file is gated to fire only on the exact buggy trigger. I ran the no-op cases — nn.Linear, standard eval BatchNorm, mean(-1) RMSNorm, scalar/finite clamp, dropout=0 SDPA, stride-1 slice — and they produce identical IR on origin and this branch. So nothing we ship is touched.

Your call on scope: since 5/6 are defensive guards for patterns no current model emits, I'm happy to slim this PR to just the lower.py inf-materialization fix (the one genuinely reachable bug) if you'd rather not carry guards you can't trip — or keep the rest as zero-regression hardening that fails loudly / falls back instead of silently miscompiling for hand-written modules and future models. Let me know which you prefer. cc @jakmro

[kar-m]

@jakmro with the context given by Noah about reproducibility I don't know if this should be merged -- deferring to you.

[ncylich]

@kar-m @jakmro following up on the reproducibility/regression concern — and correcting something I said above.

Correction: I previously called lower.py's inf-materialization "the one genuinely reachable fix." That was wrong. Its only behavioral change is on the non-FP16 branch, but the engine's working precision is FP16 and the where/clamp lowering forces operands through _ensure_fp16_tensor first, so that branch can never execute. lower.py is latent like the rest.

So I took the conservative path and removed every change that no model — current or future — can reach on the FP16 engine + eval-mode export pipeline:

• lower.py (both hunks) — FP16 working-precision invariant; the non-FP16 path is dead.
• fusion/rms_norm.py non-last-dim guard — RMSNorm normalizes the last (feature) dim in every realistic arch; the guard only fires on a shape that doesn't occur, and would needlessly skip a valid fusion when shape is unknown.
• import_semantics.py SDPA nonzero-dropout guard — model.eval() at capture forces attention_dropout → 0, so dropout_p is always 0 at export and the guard never fires.

What remains — each is reachable, and none can regress a model we support today:

File	Why it's reachable	Regression risk
aten_ops.py	canonical_ir_op matched the first prefix in declaration order, so aten.addmm→add, aten.minimum→min, aten.maximum→max, aten.slice_scatter→slice, aten.select_scatter→index were all mis-routed. Longest-prefix matching fixes it.	None — no supported model emits these bare ops (suite is green on origin with the mis-routing), so it only changes future models.
importers.py	(a) scaled addmm (beta/alpha != 1): origin silently computes bias + x@w and drops the scale → now fails closed. (b) eval-mode BatchNorm: our pipeline runs no run_decompositions, so torch.export emits _native_batch_norm_legit_no_training (7 args); origin reads args[7] → IndexError.	None — no supported model uses scaled addmm or BatchNorm; this fixes a crash/silent-wrong-output for future conv models.
import_semantics.py	A literal SDPA mask is stored in attrs["mask"] and read by lowering (lower.py:1145), but the SDPA→attention rewrite rebuilt attrs and dropped it. Now preserved — mirrors the identical passthrough already in canonicalize/cleanup.py.	None — strictly preserves a mask that was being lost; restores correct masking.
capture_jax.py	JAX path: a strided slice (stride != 1, full extent) was aliased as a no-op; plus a duplicate gather-squeeze reshape.	None — JAX-transpiler path, same engine formats.

Net: the removed changes restore origin behavior exactly for every supported model; the kept changes are a crash-fix, a fail-closed guard against silent wrong output, a mask-preservation fix, and JAX-path slice/gather correctness. Verified green: pytest python/tests (191 passed, 2 skipped), pytest python/cactus/convert/tests (53 passed). PR title/scope updated accordingly.

[ncylich]

@kar-m @jakmro I went back and actually ran each change against origin/main vs this branch on the real capture/import pipeline (instead of reasoning about it), and it corrected two of my earlier calls. Net: the branch is now 4 files, every one reached by a runnable repro, and nothing here can regress a supported model.

Reached (kept):

• aten_ops.py — normalize_target on origin: aten.addmm→add, aten.minimum→min, aten.maximum→max, aten.slice_scatter→slice, aten.select_scatter→index (all mis-routed); fixed here. A model with torch.minimum/maximum/addmm emits these — MinMax model imports to {minimum, maximum, add} on this branch.
• importers.py (scaled addmm) — torch.addmm(bias, a, b, beta=2.0) now fails closed with a clear error instead of silently dropping beta.
• fusion/rms_norm.py — this one I had wrongly proposed dropping. A plausible channel-wise norm x * rsqrt(x.pow(2).mean(1, keepdim=True) + eps) * w over (B,C,H,W) in FP16 gets fuse_rms_norm = True and the fused node's weight comes out (5,) (the last dim) while the mean was over the channel dim (4) → silently normalizes the wrong axis. With the guard: that case is left unfused (correct), and a normal last-dim RMSNorm still fuses. So I restored it.
• capture_jax.py — x[::2] on a length-6 vector: origin returns the full length-6 input aliased as a no-op (ops=[], out_shape=(6,)); this branch produces a real slice → out_shape=(3,), matching eager JAX [0,2,4].

Dropped as unreachable by any model on the FP16 engine / standard import pipeline:

• lower.py — its only behavioral change is the non-FP16 branch; the FP16 working precision never executes it.
• import_semantics.py (entire) — SDPA dropout guard never fires (model.eval() forces dropout_p=0); and the literal-mask passthrough is dead: attrs["mask"] is only set for a non-tensor literal mask, but SDPA attn_mask is always a Tensor/None (tested tensor-arg, inline-const, and is_causal forms — none set it).
• importers.py (BatchNorm no-training branch) — directly relevant to your "could this regress a supported model" worry: _native_batch_norm_legit_no_training only appears after run_decompositions, which in this repo runs only on the NPU→CoreML path (npu/vision.py, npu/audio.py → ct.convert), never feeding import_captured_to_ir. The IR importer always sees the 8-arg aten.batch_norm.default, handled identically before/after. So I removed it.

Verified green with JAX installed so the capture tests actually run: pytest python/tests + pytest python/cactus/convert/tests → 264 passed, 1 skipped.

[ncylich]

For anyone who wants to re-run the reachability checks: I committed the repro scripts in 4d4f745 and immediately reverted them in the next commit, so they're checkout-able by SHA but stay out of the squash/merge (the PR's net diff is still just the 4 fix files).

git checkout 4d4f7451ee36 -- python/tests/transpile/tools/repro_706
python python/tests/transpile/tools/repro_706/repro_pytorch.py     # routing, batchnorm form, scaled-addmm, SDPA mask, rms_norm
pip install "jax[cpu]" && python python/tests/transpile/tools/repro_706/repro_jax.py   # strided-slice aliasing
bash python/tests/transpile/tools/repro_706/origin_vs_pr.sh        # origin/main vs this branch, one file swapped at a time

README.md in that dir maps each script to the kept/dropped decision.

[Copilot]

Copilot was unable to review this pull request because the user who requested the review has reached their quota limit.