Skip to content

Add bessel_i0e and bessel_i1e ops#3193

Open
robert-johansson wants to merge 1 commit intoml-explore:mainfrom
robert-johansson:add-bessel-i0e-i1e
Open

Add bessel_i0e and bessel_i1e ops#3193
robert-johansson wants to merge 1 commit intoml-explore:mainfrom
robert-johansson:add-bessel-i0e-i1e

Conversation

@robert-johansson
Copy link
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>
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.

1 participant