Skip to content

Add bessel_i0e and bessel_i1e ops - #3193

Closed
robert-johansson wants to merge 1 commit into
ml-explore:mainfrom
robert-johansson:add-bessel-i0e-i1e
Closed

robert-johansson wants to merge 1 commit into
ml-explore:mainfrom
robert-johansson:add-bessel-i0e-i1e

Conversation

@robert-johansson

Copy link
Copy Markdown
Contributor

Summary

Adds element-wise exponentially scaled modified Bessel functions of the first kind:

  • bessel_i0e(x) = exp(-|x|) I₀(x)
  • bessel_i1e(x) = exp(-|x|) I₁(x)

These are the numerically stable building blocks for Bessel-dependent distributions (von Mises, Rice, noncentral chi-squared) in probabilistic programming and physics applications. Relates to #21 (bessel functions requested).

  • Metal: Cephes Chebyshev polynomial approximation (new kernel bessel.h)
  • CPU: Same Cephes algorithm via SIMD math helpers
  • CUDA: Same Cephes algorithm in device functors
  • Autograd: grad(i0e) = i1e(x) - sign(x)·i0e(x), grad(i1e) = i0e(x) - (sign(x) + 1/x)·i1e(x) with safe x=0 handling
  • Python bindings: mx.bessel_i0e() and mx.bessel_i1e() with docstrings

The algorithm is Cephes (public domain, Stephen Moshier) — the same implementation used by SciPy, JAX/XLA, PyTorch, and TensorFlow. Float32 relative error < 2e-7 across the full range.

Test plan

  • C++ tests: 15 assertions covering accuracy vs known values, even/odd symmetry, large x (asymptotic), type promotion, float16, vectorized shapes, and gradient correctness for both ops
  • Python tests: equivalent coverage with numpy comparison
  • CUDA tests (no local hardware — relies on CI)

🤖 Generated with Claude Code

Co-Authored-By: Claude Opus 4.6 <noreply@anthropic.com>
@zcbenz

zcbenz commented Jul 1, 2026

Copy link
Copy Markdown
Member

Sorry for the long wait, it had been hard to justify the maintenance costs of adding new primitives against the very few requests of these ops. For the moment I'm going to close this issue because there are some planned improvements on SIMD in #3019, and the discussion of adding new special functions can happen in #3330. Still thanks for the PR!

@zcbenz zcbenz closed this Jul 1, 2026
@robert-johansson
robert-johansson deleted the add-bessel-i0e-i1e branch July 27, 2026 14:11
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.

2 participants