chebax

Differentiable special functions for JAX and GPUs, built from Chebyshev approximants. A build-time generator (mpmath precision) produces fixed-degree branchless polynomial kernels with exact derivative series — including gradients with respect to function parameters (Bessel order, Beta shape, von Mises concentration), which mainstream ML libraries lack. No mpmath at use time; the runtime imports only jax and numpy.

pip install chebax

import jax
jax.config.update("jax_enable_x64", True)   # before building instances: tables
                                            # materialize in the active precision
import jax.numpy as jnp
import chebax

jv = chebax.besselj(2.5)                    # any real order in [0, 10]; cached
x = jnp.linspace(0.1, 30.0, 200)
jv(x), jax.vmap(jax.grad(jv))(x)            # values and dJ/dx, jit/vmap-safe

chebax.besselk_fn(1.7, x)                   # order as a traced scalar:
jax.grad(chebax.besselk_fn)(1.7, 5.0)       #   jax.grad w.r.t. the ORDER
chebax.betainc_fn(2.0, 3.0, 0.4)            # dI/da, dI/db via jax.grad (jax#38610)
chebax.betaincinv(2.0, 3.0, 0.05)           # differentiable Beta quantile (jax#2399)

What's in

Plus the generic core (fit, ChebSeries, PiecewiseCheb) and bake emitters (chebax.bake.jax_module, chebax.bake.xsf_header: self-contained pure-jax modules and C++17 headers from an instance).

For pymc users: import chebax.pytensor (installs with pip install chebax[pytensor]) registers JAX-backend lowerings for the pytensor ops that otherwise need tfp-nightly under pm.sample(nuts_sampler="numpyro"|"blackjax") — betaincinv, gammaincinv, gammainccinv, erfcx, erfcinv, ive, kve — and adds betainc gradients in its shape parameters, so censored or truncated StudentT/Beta likelihoods with a latent scalar shape sample instead of raising. Shape parameters must be scalar (batched ones fall back to tfp or fail loudly); out-of-domain values return nan, never silently wrong numbers.

notebooks/ holds themed, executed walkthroughs: why Chebyshev nodes work, the Bessel family with a learnable Matérn kernel, differentiable quantiles, truncated and circular distributions in numpyro, Gaussian tails / Lambert W / binomial reliability, copulas, and baking. The numpyro and plotting dependencies install with pip install chebax[examples]. examples/ holds plain scripts: a few of the notebook workflows in copy-paste form, plus integration examples with no notebook counterpart (a GIG distribution for efax, learned Matérn smoothness, truncated sampling).

Parameters come first, evaluation point last (opposite of scipy). Orders/shapes are uniform per call; per-element parameter arrays are out of scope by design. Accuracy is validated against mpmath at 40 dps — measured worst cases and the per-family error metrics are in each test file's docstring; tables regenerate bit-for-bit from checked-in generators.

Docs

docs/adding-a-recipe.md is the workflow for new functions; docs/increments.md the design log of every recipe (what was measured, which traps were found); PROJECT.md the plan and evidence; CLAUDE.md how to work in the repo. Grown out of a private research project on GPU Bessel functions; its two load-bearing measurements are reproduced here in experiments/. BSD-3-Clause.