Skip to content

Fix Metal expm1 argument reduction accuracy - #4780

Merged
zcbenz merged 2 commits into
ml-explore:mainfrom
arcusbuilds:fix/expm1-ln2-correction
Oct 31, 2026
Merged

zcbenz merged 2 commits into
ml-explore:mainfrom
arcusbuilds:fix/expm1-ln2-correction

Conversation

@arcusbuilds

Copy link
Copy Markdown
Contributor

Restore the low-part ln(2) correction in the Metal expm1f argument reduction. Omitting it biases the reduced argument and propagates accuracy errors to expm1, erf, and exact GELU.

Add a small GPU regression test comparing eight float32 inputs against NumPy float64 results. The test fails without the correction and passes with it.

Validation:

  • Source build with MLX_METAL_JIT=ON succeeded.
  • Six relevant expm1, erf, and GELU tests passed.
  • Ruff lint and format checks, stream argument checks, and git diff --check passed.

Fixes #4765

@zcbenz
zcbenz force-pushed the fix/expm1-ln2-correction branch from 3d2ddc7 to fb71a24 Compare October 31, 2026 02:13
@zcbenz
zcbenz merged commit 700520a into ml-explore:main Oct 31, 2026
24 checks passed
malfet added a commit to pytorch/pytorch that referenced this pull request Oct 31, 2026
Fixes #202520

Adds the missing low-part FMA of the `a - j*ln2` reduction to `c10::metal::expm1f` (vendored from MLX in #155611, same fix as ml-explore/mlx#4780), which makes it <1 ulp on [-88, 88] instead of up to ~2900 ulp for |x| > ln2/2, and drops the `|x| < 0.1` cutoff in `expm1_functor` that worked around it (and left ~10 ulp just above 0.1). ELU/CELU/SELU now use `c10::metal::expm1f` instead of `precise::exp(x) - 1`, which returned 0 for small negative inputs

Extends `test_expm1_small` to ~79 and the ln2/2 boundary with rtol=3e-7 (fails on MPS before this change) and removes the now-passing complex64 `test_reference_numerics_small` xfail for expm1. Adds a similar `test_elu_small` for ELU/CELU/SELU (also fails on MPS before this change).

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
ghstack-source-id: 47552eb
Pull-Request: #202897
pytorchmergebot pushed a commit to pytorch/pytorch that referenced this pull request Nov 1, 2026
Fixes #202520

Adds the missing low-part FMA of the `a - j*ln2` reduction to `c10::metal::expm1f` (vendored from MLX in #155611, same fix as ml-explore/mlx#4780), which makes it <1 ulp on [-88, 88] instead of up to ~2900 ulp for |x| > ln2/2, and drops the `|x| < 0.1` cutoff in `expm1_functor` that worked around it (and left ~10 ulp just above 0.1). ELU/CELU/SELU now use `c10::metal::expm1f` instead of `precise::exp(x) - 1`, which returned 0 for small negative inputs

Extends `test_expm1_small` to ~79 and the ln2/2 boundary with rtol=3e-7 (fails on MPS before this change) and removes the now-passing complex64 `test_reference_numerics_small` xfail for expm1. Adds a similar `test_elu_small` for ELU/CELU/SELU (also fails on MPS before this change).

🤖 Generated with [Claude Code](https://claude.com/claude-code)

Co-Authored-By: Claude Opus 5.5 <noreply@anthropic.com>
Pull Request resolved: #202897
Approved by: https://github.com/Skylion007
ghstack dependencies: #202890
Sign up for free to join this conversation on GitHub. Already have an account? Sign in to comment

Labels

None yet

Projects

None yet

Development

Successfully merging this pull request may close these issues.

[BUG] Metal expm1 (and erf for |x| > 0.93) reads up to 336 ulp high: the argument reduction drops ln 2's low part

2 participants