REVIEW 7 cited by
Equinox: neural networks in JAX via callable PyTrees and filtered transformations
Not yet reviewed by Pith; the record is open.
This paper has not been read by Pith yet. Machine review is queued; the pith claim, tier, and objections will appear here once it completes.
SPECIMEN: schema-true, not a live event
T0 review · schema-true
One-sentence machine reading of the paper's core claim.
pith:XXXXXXXX · record.json · timestamp
read the original abstract
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 \url{https://github.com/patrick-kidger/equinox}.
Forward citations
Cited by 7 Pith papers
-
A Fortran General-Purpose Transpiler: Proof of Concept
FGPT deterministically translates isolated Fortran procedures into numerically verified NumPy and JAX/Equinox modules, with forward outputs matching Fortran to ~1e-9, but explicit auto-differentiation is not yet implemented.
-
Neural Network-Augmented Pfaffian Wave-functions for Scalable Simulations of Interacting Fermions
A neural-network Pfaffian wavefunction with hidden fermions achieves state-of-the-art variational accuracy on the 2D Hubbard model and scales via low-rank updates and sublattice symmetry.
-
Incorporating Hierarchical Semantics in Sparse Autoencoder Architectures
A two-level mixture-of-experts sparse autoencoder models parent and child concepts together, improving reconstruction and reducing feature redundancy on Gemma 2-2B activations compared to flat top-k SAEs.
-
Fast, accurate, and differentiable: a neural-network surrogate for NRSur7dq4 precessing binary black hole waveforms
A piecewise MLP surrogate emulates NRSur7dq4 over its full domain at NR-faithful accuracy with ~1 ms GPU latency and a fully differentiable JAX likelihood pipeline.
-
Differentiable Cardiac Electrophysiology Simulations for Dynamical State and Parameter Estimation
A differentiable solver for cardiac reaction-diffusion waves recovers hidden states and parameters of the Aliev-Panfilov model from sparse, surface-only, or noisy voltage observations, including two experimental monol...
-
Biomechanics-aware Multi-view Markerless Motion Capture of Dexterous Hand Movements
End-to-end biomechanics-aware optimization recovers plausible multi-joint hand kinematics from multi-view markerless video more robustly than two-stage triangulation-plus-IK, especially under object occlusion.
-
Uncertainty in Physics and AI: Taxonomy, Quantification, and Validation
A unified taxonomy of uncertainty in ML for physics is introduced together with validation tools such as coverage, calibration, and proper scoring rules, illustrated on regression and classification tasks.
Discussion (0). Continue with ORCID to comment.