Equinox Python Packages

equinox

Elegant easy-to-use neural networks + scientific computing in JAX. https://docs.kidger.site/equinox/

lineax

Linear solvers in JAX and Equinox. https://docs.kidger.site/lineax

optimistix

Nonlinear optimisation (root-finding, least squares, ...) in JAX+Equinox. https://docs.kidger.site/optimistix/

diffrax

Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/

exponax

Efficient Differentiable n-d PDE Solvers in JAX.

pdequinox

Neural Emulator Architectures in JAX.

trainax

Neural Emulator Architectures in JAX.

quax

Multiple dispatch over abstract array types in JAX.

diffraxtra

Extras for Diffrax: OOP and vectorization

oncequinox

micro-library for making Equinox singletons

eqx-zoo

Verified Equinox ports of pretrained models, numerically matched against Hugging Face

esm2quinox

An implementation of ESM2 in Equinox+JAX

xtrax

Composable building blocks for JAX/Equinox training loops (alpha, experimental)

aminx

JAX/Equinox implementation of ProteinMPNN and LigandMPNN for protein sequence design

parametrix

🪄 Flax-like computed parameters for bare JAX (and Equinox)

mimosa-ml

MIMOSA: a fully-featured multi-task Gaussian process framework for analysing functional data.

eqxbatch

Batched evaluation of Equinox modules: a vmap that lives in the PyTree.

eqxvision

A Python package of computer vision models for the Equinox ecosystem.

sbgm

Score-based Diffusion models in JAX.

frontx

⚛️ Nonlinear diffusion problems with JAX

jax-blox

A functional and lightweight neural network library for JAX.

jaxon

Save pytrees efficiently in hdf5 files.

lerax

Reinforcement Learning in Jax

finslerax

Differentiable Finsler geometry in JAX: learnable asymmetric metrics, geodesics, eikonal solvers, curvature, and parallel transport.