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

Patrick Kidger, Cristian Garcia

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-o…