mlsanityData to 7 Oct 2026

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

arXiv 2111.00254 ↗First posted 30 Oct 2021cs.LG cs.PLcited by 143

Not on the map. Why?

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.

Key facts. Cited by 143 papers in the map (data through 7 Oct 2026). Builds on 4 of its 12 references that are in the map. First posted on arXiv 30 Oct 2021.

Builds on

4 of 12 references.

References · 4

  1. 2019PyTorch: An Imperative Style, High-Performance Deep Learning Librarycited by 20,160
  2. 2018Fashionable Modelling with Fluxcited by 53
  3. 2021Swift for TensorFlow: A portable, flexible platform for deep learningcited by 6
  4. 2021Decomposing reverse-mode automatic differentiationcited by 4

Built on by

Cited by year (2026 to 7 Oct)

Citing papers per year: 2022: 3; 2023: 12; 2024: 41; 2025: 50; 2026: 37.

143 papers.

Sort:

By area:

  1. 2024Adaptive Data Optimization: Dynamic Sample Selection with Scaling Lawscited by 46
  2. 2022Robust Neural Posterior Estimation and Statistical Model Criticismcited by 45
  3. 2023Parallelizing non-linear sequential models over the sequence lengthcited by 45
  4. 2023Disentanglement via Latent Quantizationcited by 40
  5. 2024Benchmarking Predictive Coding Networks -- Made Simplecited by 29
  6. 2023GECCO: Geometrically-Conditioned Point Diffusion Modelscited by 24
  7. 2024Scaling physics-informed hard constraints with mixture-of-expertscited by 24
  8. 2025On the Edge of Memorization in Diffusion Modelscited by 22
  9. 2024Differentiable Biomechanics Unlocks Opportunities for Markerless Motion Capturecited by 18
  10. 2024Efficient mapping of phase diagrams with conditional Boltzmann Generatorscited by 18
  11. 2025FlowQ: Energy-Guided Flow Policies for Offline Reinforcement Learningcited by 14
  12. 2024Tripod: Three Complementary Inductive Biases for Disentangled Representation Learningcited by 13
38 more
  1. 2024Learning from Linear Algebra: A Graph Neural Network Approach to Preconditioner Design for Conjugate Gradient Solverscited by 13
  2. 2025Neural ODE Transformers: Analyzing Internal Dynamics and Adaptive Fine-tuningcited by 12
  3. 2023XLB: A differentiable massively parallel lattice Boltzmann library in Pythoncited by 11
  4. 2024Neural Context Flows for Meta-Learning of Dynamical Systemscited by 9
  5. 2025Residual Matrix Transformers: Scaling the Size of the Residual Streamcited by 8
  6. 2025Portable Biomechanics Laboratory: Clinically Accessible Movement Analysis from a Handheld Smartphonecited by 8
  7. 2024Align Your Intents: Offline Imitation Learning via Optimal Transportcited by 7
  8. 2024Fusing uncalibrated IMUs and handheld smartphone video to reconstruct knee kinematicscited by 6
  9. 2024NeuroSEM: A hybrid framework for simulating multiphysics problems by coupling PINNs and spectral elementscited by 6
  10. 2025Training a Foundation Model for Materials on a Budgetcited by 6
  11. 2023Self-Supervised Learning of Gait-Based Biomarkerscited by 5
  12. 2024Penzai + Treescope: A Toolkit for Interpreting, Visualizing, and Editing Models As Datacited by 5
  13. 2025Towards Foundational Models for Dynamical System Reconstruction: Hierarchical Meta-Learning via Mixture of Expertscited by 5
  14. 2025Neural Flow Samplers with Shortcut Modelscited by 5
  15. 2025Scaling sparse feature circuit finding for in-context learningcited by 5
  16. 2025Weight-Space Linear Recurrent Neural Networkscited by 5
  17. 2025Differentiable Simulation of Hard Contacts with Soft Gradients for Learning and Controlcited by 5
  18. 2025JAX-MPM: A Learning-Augmented Differentiable Meshfree Framework for GPU-Accelerated Lagrangian Simulation and Geophysical Inverse Modelingcited by 5
  19. 2026Monocular Biomechanical Tracking of Fingers with Inverse Kinematics to Foundation Modelscited by 5
  20. 2023Noncommutative C^*-algebra Net: Learning Neural Networks with Powerful Product Structure in C^*-algebracited by 4
  21. 2023Hybridizing Physics and Neural ODEs for Predicting Plasma Inductance Dynamics in Tokamak Fusion Reactorscited by 4
  22. 2024textttimmrax: A Parallelizable and Differentiable Toolbox for Interval Analysis and Mixed Monotone Reachability in JAXcited by 4
  23. 2024Engineering morphogenesis of cell clusters with differentiable programmingcited by 4
  24. 2024Separable Operator Networkscited by 4
  25. 2025Scaling Up Resonate-and-Fire Networks for Fast Deep Learningcited by 4
  26. 2025A Simultaneous Approach for Training Neural Differential-Algebraic Systems of Equationscited by 4
  27. 2025A Unifying Framework for Parallelizing Sequential Models with Linear Dynamical Systemscited by 4
  28. 2024Certified Robust Invariant Polytope Training in Neural Controlled ODEscited by 3
  29. 2024Consistent machine learning for topology optimization with microstructure-dependent neural network material modelscited by 3
  30. 2024SNNAX -- Spiking Neural Networks in JAXcited by 3
  31. 2024Constrained Latent Action Policies for Model-Based Offline Reinforcement Learningcited by 3
  32. 2025Biomechanical Reconstruction with Confidence Intervals from Multiview Markerless Motion Capturecited by 3
  33. 2025PRDP: Progressively Refined Differentiable Physicscited by 3
  34. 2025Permutation Equivariant Neural Controlled Differential Equations for Dynamic Graph Representation Learningcited by 3
  35. 2025Score-based generative emulation of impact-relevant Earth system model outputscited by 3
  36. 2026MapPFN: Learning Causal Perturbation Maps in Contextcited by 3
  37. 2026Uncertainty in Physics and AI: Taxonomy, Quantification, and Validationcited by 3
  38. 2023Data-driven multiscale modeling for correcting dynamical systemscited by 2

50 most cited of 143 All 143 (JSON)