Equinox

Written by Patrick Kidger, Equinox fills the gaps around core JAX without wrapping it in a separate framework. A model is an ordinary Python class that subclasses eqx.Module, which registers it as a PyTree, so the same object can be handed straight to jax.jit, jax.grad, or vmap. Filtered transformations such as filter_jit and filter_grad decide which leaves are arrays to trace and which are static Python values, sparing much of the bookkeeping Flax or Haiku users do by hand. The eqx.nn package covers standard layers from linear and convolution to attention, recurrent cells, and stateful batch norm, while other modules add tree surgery with tree_at, runtime checks via error_if that survive compilation, and serialisation helpers. Worked examples include BERT, a vision transformer, and score-based diffusion. Installation is pip install equinox on Python 3.10 or newer, with hardware support inherited from JAX on CPU, GPU, or TPU. Apache-2.0 makes commercial use straightforward. It suits researchers and scientific computing users who like JAX but want PyTorch-style model code; the author's Diffrax builds on it.