JAX Stack

Here are my favourite JAX packages.

Tutorials¶

• Jax Tutorial 101
• PyTrees | Jax PyTrees Tutorial
• Differentiable Programming
• Autodidax: Jax From Scratch

Special DataStructures¶

• coordinax
• jaxdf
• tree-math
• xarray_jax
• quax
• jaxtyping

Symbolic Math¶

• sympy2jax

Numerical Methods¶

Linear Algebra¶

• Cola, Lab
• autoray
• einx
• einops
• matfree - Matrix-Free linear algebra in JAX
• opt-einsum - optimized einsum (numpy, JAX, TF, PyTorch, Dask, CuPy, Sparse)

Convolutions¶

• Kernex
• Serket

Integration¶

torchquad,

Interpolation¶

• interpax
• Nyx
• RBF

Optimization¶

• Optimistix
• LineaX
• Optax
• JaxOp
• varz - Simple, multi-backend constrained (L-BFGS) and unconstrained optimization (Adam).
• ott

Kernels¶

• mlkernels - Kernel Matrices (JAX, TF, PyTorch, Julia).
• KernelBiome

Differentiation¶

• FiniteDiffX, FiniteVolX, SpectralDiffX
• jax-fem
• Probfindiff
• LapJax
• RBF-FDax

ODESolvers¶

• Diffrax,
• probdiffeq - probabilistic solvers for differential equations

Ordinary Differential Equations¶

• DiffEqZoo
• Dysts
• sdeint

Partial Differential Equations¶

• pyqg-jax pyshocks
• somaX,
• dinosaur
• veros

Basis Functions¶

PCA/SVD/POD/EOF¶

• pcax

Fourier¶

• s2fft

Orthogonal¶

• orthojax
• orthax
• s2ball

Wavelet¶

• jax-wavelet-toolbox
• cr-wavelets
• s2wav

Spherical Harmonics¶

• SphericalHarmonics - spherical harmonics (numpy, JAX, PyTorch, TF)
• Jax Implementation

Neural Networks¶

• Equinox
• Flax
• Keras
• serket

Probabilistic¶

Probabilistic Programming Languages¶

• blackjack,
• numpyro
• numpyro-ext
• tfp.substrate.jax
• bayeux

Distributions¶

• efax
• fenbux

Samplers¶

• jaxns
• SGMCMCJax - stochastic Gradient samplers in jax
• jax-sgmc - stochastic gradient samplers in jax

Normalizing Flows¶

• FlowMC
• flowjax

Gaussian Processes¶

• GPJax
• TinyGP

State Space Models¶

• Dynamax
• ReBayes
• sts-jax
• dynax

Image Processing¶

• dm-pix

Units¶

• jpu

Parallel Programming¶

• mpi4jax
• paxml
• Jax-Parallel