# Equinox: neural networks in JAX via callable PyTrees and filtered transformations

arXiv: https://arxiv.org/abs/2111.00254 · HTML: https://mlsanity.com/paper/2111.00254/

## Key facts

- Cited by 148 papers in the map (data through 14 Oct 2026).
- Builds on 4 of its 12 references that are in the map.
- First posted on arXiv 30 Oct 2021.

## Builds on (most cited first, 4 of 4)

- [PyTorch: An Imperative Style, High-Performance Deep Learning Library](https://mlsanity.com/paper/1912.01703/) (2019)
- [Fashionable Modelling with Flux](https://mlsanity.com/paper/1811.01457/) (2018)
- [Swift for TensorFlow: A portable, flexible platform for deep learning](https://mlsanity.com/paper/2102.13243/) (2021)
- [Decomposing reverse-mode automatic differentiation](https://mlsanity.com/paper/2105.09469/) (2021)

## Built on by (most cited first, 30 of 148)

- [Adaptive Data Optimization: Dynamic Sample Selection with Scaling Laws](https://mlsanity.com/paper/2410.11820/) (2024)
- [Robust Neural Posterior Estimation and Statistical Model Criticism](https://mlsanity.com/paper/2210.06564/) (2022)
- [Parallelizing non-linear sequential models over the sequence length](https://mlsanity.com/paper/2309.12252/) (2023)
- [Disentanglement via Latent Quantization](https://mlsanity.com/paper/2305.18378/) (2023)
- [Benchmarking Predictive Coding Networks -- Made Simple](https://mlsanity.com/paper/2407.01163/) (2024)
- [GECCO: Geometrically-Conditioned Point Diffusion Models](https://mlsanity.com/paper/2303.05916/) (2023)
- [Scaling physics-informed hard constraints with mixture-of-experts](https://mlsanity.com/paper/2402.13412/) (2024)
- [On the Edge of Memorization in Diffusion Models](https://mlsanity.com/paper/2508.17689/) (2025)
- [Differentiable Biomechanics Unlocks Opportunities for Markerless Motion Capture](https://mlsanity.com/paper/2402.17192/) (2024)
- [Efficient mapping of phase diagrams with conditional Boltzmann Generators](https://mlsanity.com/paper/2406.12378/) (2024)
- [Learning from Linear Algebra: A Graph Neural Network Approach to Preconditioner Design for Conjugate Gradient Solvers](https://mlsanity.com/paper/2405.15557/) (2024)
- [FlowQ: Energy-Guided Flow Policies for Offline Reinforcement Learning](https://mlsanity.com/paper/2505.14139/) (2025)
- [Tripod: Three Complementary Inductive Biases for Disentangled Representation Learning](https://mlsanity.com/paper/2404.10282/) (2024)
- [Neural ODE Transformers: Analyzing Internal Dynamics and Adaptive Fine-tuning](https://mlsanity.com/paper/2503.01329/) (2025)
- [Residual Matrix Transformers: Scaling the Size of the Residual Stream](https://mlsanity.com/paper/2506.22696/) (2025)
- [XLB: A differentiable massively parallel lattice Boltzmann library in Python](https://mlsanity.com/paper/2311.16080/) (2023)
- [Neural Context Flows for Meta-Learning of Dynamical Systems](https://mlsanity.com/paper/2405.02154/) (2024)
- [Align Your Intents: Offline Imitation Learning via Optimal Transport](https://mlsanity.com/paper/2402.13037/) (2024)
- [Portable Biomechanics Laboratory: Clinically Accessible Movement Analysis from a Handheld Smartphone](https://mlsanity.com/paper/2507.08268/) (2025)
- [Fusing uncalibrated IMUs and handheld smartphone video to reconstruct knee kinematics](https://mlsanity.com/paper/2405.17368/) (2024)
- [NeuroSEM: A hybrid framework for simulating multiphysics problems by coupling PINNs and spectral elements](https://mlsanity.com/paper/2407.21217/) (2024)
- [Training a Foundation Model for Materials on a Budget](https://mlsanity.com/paper/2508.16067/) (2025)
- [Self-Supervised Learning of Gait-Based Biomarkers](https://mlsanity.com/paper/2307.16321/) (2023)
- [Separable Operator Networks](https://mlsanity.com/paper/2407.11253/) (2024)
- [Penzai + Treescope: A Toolkit for Interpreting, Visualizing, and Editing Models As Data](https://mlsanity.com/paper/2408.00211/) (2024)
- [Towards Foundational Models for Dynamical System Reconstruction: Hierarchical Meta-Learning via Mixture of Experts](https://mlsanity.com/paper/2502.05335/) (2025)
- [Neural Flow Samplers with Shortcut Models](https://mlsanity.com/paper/2502.07337/) (2025)
- [Scaling sparse feature circuit finding for in-context learning](https://mlsanity.com/paper/2504.13756/) (2025)
- [Weight-Space Linear Recurrent Neural Networks](https://mlsanity.com/paper/2506.01153/) (2025)
- [Differentiable Simulation of Hard Contacts with Soft Gradients for Learning and Control](https://mlsanity.com/paper/2506.14186/) (2025)

## Abstract (arXiv)

JAX and PyTorch are two popular Python autodifferentiation frameworks. JAX is based around pure functions and functional programming. PyTorch has popularised the use of an object-oriented (OO) class-based syntax for defining parameterised functions, such as neural networks. That this seems like a fundamental difference means current libraries for building parameterised functions in JAX have either rejected the OO approach entirely (Stax) or have introduced OO-to-functional transformations, multiple new abstractions, and been limited in the extent to which they integrate with JAX (Flax, Haiku, Objax). Either way this OO/functional difference has been a source of tension. Here, we introduce `Equinox', a small neural network library showing how a PyTorch-like class-based approach may be admitted without sacrificing JAX-like functional programming. We provide two main ideas. One: parameterised functions are themselves represented as `PyTrees', which means that the parameterisation of a function is transparent to the JAX framework. Two: we filter a PyTree to isolate just those components that should be treated when transforming (`jit', `grad' or `vmap'-ing) a higher-order function of a parameterised function -- such as a loss function applied to a model. Overall Equinox resolves the above tension without introducing any new programmatic abstractions: only PyTrees and transformations, just as with regular JAX. Equinox is available at urlhttps://github.com/patrick-kidger/equinox.