jaxellip

What it is and what it does

jaxellip fills a gap in JAX's scientific computing stack by providing differentiable, JIT-compatible implementations of three complete elliptic integrals: ellipk, ellipkm1, and ellipe. These functions are available in scipy.special but not in jax.scipy.special, making them incompatible with JAX's autodiff and compilation. The package uses Cephes minimax polynomials (the same method as SciPy) for forward evaluation and custom JVP rules for derivatives, enabling full differentiability in forward mode, reverse mode, and higher-order derivatives like jax.hessian.

The implementation handles numerical edge cases consistently with SciPy—ellipk(1) returns inf, ellipe(1) returns 1, and out-of-domain inputs return NaN. Results are tested against arbitrary-precision mpmath references and validated against scipy.special to about 1e-14 relative accuracy. The package requires only jax as a runtime dependency and supports Python 3.12–3.14.

Use it for:

• Compute elliptic integrals inside JAX-compiled functions (jit, vmap, grad) without leaving the JAX ecosystem.
• Differentiate functions that depend on elliptic integrals using jax.grad, jax.jacobian, or jax.hessian.
• Replace scipy.special calls in numerical code when JIT compilation or autodiff is needed.
• Implement physics or engineering models (e.g., pendulum dynamics, elliptic orbits) with automatic differentiation.
• Build machine learning models that incorporate elliptic integrals as part of a larger differentiable computation graph.

Worth the install?

AI-flagged interpretation of the facts on this page — verify before relying

Provides JAX-compatible implementations of elliptic integrals (ellipk, ellipkm1, ellipe) that are differentiable and JIT-compilable, addressing the gap in jax.scipy.special.

Yes, if you need elliptic integrals inside JAX-compiled or autodiff code. The package is actively maintained, has low install friction, carries no known vulnerabilities, and is MIT-licensed. The 0.1.0 release is recent and tested against mpmath and scipy.special. Install only if you actually use elliptic integrals; otherwise it adds an unnecessary dependency.

Install

jaxellip on PyPI

pip install jaxellip

uv add jaxellip

poetry add jaxellip

Installing jaxellip

Before you install

Low friction: pure Python wheel with a single runtime dependency on jax. Active maintenance with a recent release (44 days old) and current Python 3.12–3.14 support.

License in practice

MIT license permits unrestricted use, modification, and distribution with minimal attribution requirements.

Quickstart

pip install jaxellip

import jaxellip
import jax.numpy as jnp

input_vals = jnp.linspace(-10, 10, 1000)
jaxellip.ellipk(input_vals)      # Complete elliptic integral of the first kind
jaxellip.ellipkm1(input_vals)    # Complete elliptic integral of the first kind around m=1
jaxellip.ellipe(input_vals)      # Complete elliptic integral of the second kind

Verify before relying

• Whether the package is suitable for production use beyond the early 0.1.0 release stage.
• Real-world performance characteristics on different hardware (GPU, TPU) and JAX/XLA versions.

Package facts

Evidence: jaxellip-0.1.0-py3-none-any.whl

Tags

More Mathematics packages