JAX support

[crusaderky]

Blockers for JAX support, as of JAX 0.5.2:

Sobol

jnp.bitwise_xor.accumulate is quite slow on CPU and extremely slow on GPU. Alternative, specialized implementations exist; however not being able to query the device inside jax.jit means that one can't write them into pyscenarios without resorting to cumbersome parameters or module-level config flags.

• accumulate poor performance jax-ml/jax#28097
• Missing .device attribute inside @jax.jit jax-ml/jax#26000 (comment)

T* copula

Missing distribution functions in jax.scipy:

• jax.scipy.stats.chi2.ppf (one could use jax.random.chisquare on JAX own's PRNG, but not on Sobol)
• jax.scipy.stats.t.cdf - read below
• jax.scipy.stats.norm.ppf
• jax.scipy.special.betainc is extremely slow special.betainc, stdtr, and stdtrit are very slow jax-ml/jax#28547

[crusaderky]

• scipy.stats.chi2.ppf(x, df)
• is the same as scipy.stats.gamma.ppf(x, a=df/2, loc=0, scale=2)
• which is the same as 2 * scipy.special.gammaincinv(df/2, x)
• which is the same as 2 * cephes::igami(df/2, x) https://github.com/scipy/xsf/blob/4fff9b2cb2b5c31a0cf0b0f609d2699a5eeac53b/include/xsf/cephes/igami.h

[crusaderky]

Upstream trackers

• [FR] Support for scipy.special.gammaincinv jax-ml/jax#5350
• Add missing distribution cdf and ppf functions to scipy.stats jax-ml/jax#13291

[crusaderky]

jax.scipy.t.cdf(x, df) does not exist at the moment of writing; however it's equivalent to scipy.special.stdtr(df, x) in scipy 1.16 (TBC scipy 1.15). Note this is not in jax.scipy but fully works with JAX arrays on GPU too.

It's also trivial to backport to avoid using scipy at all.
From scipy.special._support_alternative_backends:

def stdtr(df, t):
    x = df / (t ** 2 + df)
    tail = betainc(df / 2, 0.5, x) / 2
    return xp.where(t < 0, tail, 1 - tail)

[crusaderky]

Some early benchmarks:

• 2**17-1 (131071) samples
• 500 dimensions
• AMD Ryzen 9 3950X 16-Core
• NVidia RTX 3080
• warm JIT in all cases

Backend	Sobol	Gaussian Copula(Sobol)	Student T Copula(Sobol)
Numpy+Numba	96ms	5.6s	29.1s
Dask + Numba1	5.6s	8.6s	5.41s
JAX CPU	184ms	474ms	1.9s2
JAX GPU	6ms	219ms	1.65s2,3

¹ 16 chunks along the samples axis, threads scheduler
² gammaincinv and betainc implemented by wrapping the scipy C implementation with jax.pure_callback and then parallelising it with a thread pool
³Because of (2), there are substantial transfers between CPU and GPU and an extended period during which the GPU is idle

[crusaderky]

could betainc be reimplemented with quadax?

[EDIT] Yes, but it's even slower