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

1 code implementation • 30 Oct 2021 • Patrick Kidger, Cristian Garcia

One: parameterised functions are themselves represented as `PyTrees', which means that the parameterisation of a function is transparent to the JAX framework.