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

Built on

“Julia: A fresh approach to numerical computing”
Jeff Bezanson, Alan Edelman, Stefan Karpinski and Viral Shah · 2017
Earlier work this paper cites.
“JAX: composable transformations of Python+NumPy programs”, 2018
James Bradbury et al · 2018
Earlier work this paper cites.
“Fashionable Modelling with Flux”
Original
Michael Innes et al · 2018
Earlier work this paper cites.
“Flux: Elegant Machine Learning with Julia”
Mike Innes · 2018
Earlier work this paper cites.
“PyTorch: An Imperative Style, High-Performance Deep Learning Library”
A. Paszke et al · 2019
Earlier work this paper cites.
“Chex: Testing made fun, in JAX!”, 2020
David Budden et al · 2020
Earlier work this paper cites.

“Flax: A neural network library and ecosystem for JAX”, 2020
Jonathan Heek et al · 2020
Cited alongside, same era.
“Haiku: Sonnet for JAX”, 2020
Tom Hennigan, Trevor Cai, Tamara Norman and Igor Babuschkin · 2020
Cited alongside, same era.
“Objax”, 2020
Objax Developers · 2020
Cited alongside, same era.
“Decomposing reverse-mode automatic differentiation”
Roy Frostig et al · 2021
Cited alongside, same era.
“Treex”, Accessed 2021
Cristian Garcia · 2021
Cited alongside, same era.
“functorch: JAX-like composable function transforms for PyTorch”, 2021
Horace He and Richard Zou · 2021
Cited alongside, same era.

“ torchtyping
P. Kidger · 2021
Closest in time.
“Opax”, Accessed 2021
NTT123 · 2021
Closest in time.
“Pax”, Accessed 2021
NTT123 · 2021
Closest in time.
“Swift for TensorFlow: A portable, flexible platform for deep learning”
Original
Brennan Saeta et al · 2021
Closest in time.
URL: https://jax.readthedocs.io/en/latest/jax.experimental.stax.html
“Stax”, Accessed 2021 · 2021
Closest in time.