Fix rms_norm eps inflation from max-|x| rescale; guard all-zero input

[nschlaepfer]

Fixes #2821.

Problem

The torch rms_norm handler rescales the input by m = max(|x|) to prevent fp16 overflow on ANE, but adds eps after the rescale:

sqrt(mean((x/m)^2) + eps) * m  ==  sqrt(mean(x^2) + eps * m^2)

The effective epsilon is inflated by m^2. For spiky activations (max|x| >> rms(x) — common in transformer residual streams and register/embedding vectors) this produces percent-level output errors at any compute precision, including FLOAT32 / CPU_ONLY. The rescale also introduces a 0/0 → NaN on all-zero rows, which the epsilon was originally there to prevent.

Fix

Two changes inside the handler, preserving the ANE overflow protection:

1. Rescale eps by 1/m^2 so the emitted math is algebraically identical to sqrt(mean(x^2) + eps) — the rescale becomes exact up to floating-point rounding.
2. Clamp the scale to >= 1 (mb.maximum(m, 1.0)): inputs with max|x| <= 1 cannot overflow fp16 when squared, so no rescale is needed there, and the clamp fixes the all-zero-row NaN.

The clamp also keeps the rescaled epsilon term well-behaved in fp16: whenever m is the true max, mean((x/m)^2) >= 1/N, so the (possibly underflowing) eps/m^2 term is only ever needed when m is clamped to 1 — where it equals eps exactly.

Verification (M3 Max, torch 2.11, python 3.11)

Case	Before	After
Spiky input (0.01·randn, one 25.0 component), fp32/CPU_ONLY, rel-L2 vs PyTorch	3.2e-4	3.0e-8
Standard randn, fp32, rel-L2	~1e-7	6.4e-8
Large activations (300·randn), FLOAT16 — overflow protection	finite, 4e-4	finite, 4.4e-4
All-zero input through the converted model	NaN	finite, matches PyTorch
Real 380M-param EEG transformer (QK-RMSNorm + register tokens), encoder rel-L2 at fp32	4–6%	~1e-6 (with equivalent formula)
Existing test_rms_norm.py suite	10/11 pass¹	10/11 pass¹

¹ test_dynamic_shapes crashes identically before and after this change on my machine (EnumeratedShapes predict crash on a macOS beta) — unrelated to this fix.

Both new regression tests fail on current main's translation (spiky: rel_l2=3.19e-04 vs 1e-5 gate; zero-input: NaN) and pass with this change.

Tests added

• test_spiky_input_eps_not_inflated — spiky-input parity at fp32 with a tight (1e-5) rel-L2 gate. The existing suite's rtol=1e-2 tolerance is exactly why this bug was invisible to it.
• test_zero_input_converted_model — all-zero input through the converted model (the existing test_edge_cases only exercised the PyTorch reference, not the conversion).

🤖 Generated with Claude Code