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; dtype="float" emits an f32 CUDA-shaped kernel body, and truncate_tol=1e-7 drops the converged f64 tail — measured 1.4–2.2x faster f32 evaluation, experiments/12).

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. The forward CDFs gammainc, gammaincc and betainc take their VALUES off the tables too, which is where the loop-free evaluation shows up for censored and truncated models; outside the table box they fall back to jax's own, since a forward function has to answer everywhere. Shape parameters must be scalar (batched ones fall back to tfp or fail loudly); out-of-domain values return nan, never silently wrong numbers. gammainccinv takes its upper-tail probability directly (it used to be solved at 1 - p, which collapsed everything below ~1e-17 onto one point), so both tails resolve to the same depth.

For numpyro users: from chebax.numpyro import TruncatedGamma, TruncatedBeta, TruncatedStudentT (installs with pip install chebax[numpyro]) gives the truncated distributions numpyro's location-scale machinery cannot cover (numpyro#969, numpyro#1365): full Distribution classes with reparameterized inverse-CDF sampling (has_rsample), truncation-normalized log_prob, and gradients through every parameter including the shapes, so NUTS and SVI work with a latent concentration or df. The truncation is normalized in log space, so an interval whose two ends land in the same tail still works: Gamma(3) truncated to [50, 51] has mass ~3e-20 and returns an ordinary density.

chebax.numpyro.Tweedie(mu, phi, power) is the compound Poisson-gamma law for data with exact zeros and positive amounts in one likelihood (claims, rainfall, biomass). Its log_prob is chebax.tweedie_logpdf, differentiable in the power too, so NUTS can sample p instead of taking it fixed: on 1,500 draws at (mu, phi, p) = (3, 0.6, 1.35) the posterior is p = 1.41 +- 0.05. Shape parameters are uniform per call; domain boxes are in the module docstring and checked at construction.

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, baking, and per-group shape parameters in a hierarchical model under NUTS. 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, chebax.numpyro with latent shape parameters under NUTS).

Parameters come first, evaluation point last (opposite of scipy). Orders/shapes are uniform per call, meaning one value whatever its shape: a scalar or any size-1 array, which is what numpyro and pytensor hand back after broadcasting. Per-element parameter arrays are out of scope by design and say so by name. The middle ground is chebax.pergroup(fn, group_idx): a static integer array assigns each element to a group, each group gets its own (traceable) parameter set — one group per chain, mixture component, or plate level. 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

Architecture defines module boundaries, numerical contracts, and validation; Coding style defines naming, formatting, and documentation conventions. Adding a recipe is the workflow for new functions. The October review records the latest review, how each finding was resolved, and what is open.

docs/increments.md is the design log of every recipe (what was measured, which traps were found); PROJECT.md holds the plan and evidence; CLAUDE.md has agent operating instructions and local test commands. Grown out of a private research project on GPU Bessel functions; its two load-bearing measurements are reproduced here in experiments/. BSD-3-Clause.