spexial

scipy.special in JAX

spexial provides special functions for JAX, following the scipy.special API. The implementations are written in terms of JAX primitives, so they compose with jit, grad and vmap, and run on CPU, GPU and TPU.

It exists to fill gaps in jax.scipy.special: zeta accepts negative integers where JAX returns nan, spence accepts complex arguments where JAX is real-only, and the modified Bessel K, Gegenbauer and polylogarithm functions have no JAX counterpart at any version. Where SciPy already covers a case, spexial matches it rather than claiming to exceed it.

Installation

pip install spexial

or

uv add spexial

Example

>>> import jax
>>> jax.config.update("jax_enable_x64", True)

>>> import jax.numpy as jnp
>>> import spexial as sp

>>> # Gegenbauer polynomial C_n^alpha(x), matching scipy.special.eval_gegenbauer
>>> sp.eval_gegenbauer(3, 0.5, 0.25)
Array(-0.3359375, dtype=float64, weak_type=True)

Everything is vectorisable and differentiable in the usual way:

>>> xs = jnp.linspace(-1.0, 1.0, 5)
>>> jax.vmap(lambda x: sp.eval_gegenbauer(3, 0.5, x))(xs)
Array([-1.    ,  0.4375, -0.    , -0.4375,  1.    ], dtype=float64)

What is here

Read Accuracy and domains before relying on any of these. It records, per function, the domain each is tested over and the tolerance it actually meets. Some are not machine-precision — the modified Bessel functions are accurate to about 1e-7, and zeta does not implement the critical strip.

Documentation

https://jaxtronomy.github.io/spexial

Development

We welcome contributions! See CONTRIBUTING.md.

uv sync --group dev
uv run nox -s all      # lint -> test -> docs

Citation

If you use spexial in work you publish, please cite it — see CITATION.cff. Several routines originate in the LINX code.

License

MIT. See LICENSE.