1 Hit in 2.8 sec

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

Patrick Kidger, Cristian Garcia
Overall Equinox resolves the above tension without introducing any new programmatic abstractions: only PyTrees and transformations, just as with regular JAX.  ...  transformations, multiple new abstractions, and been limited in the extent to which they integrate with JAX (Flax, Haiku, Objax).  ...  Acknowledgments and Disclosure of Funding PK was supported by the EPSRC grant EP/L015811/1 and by the Alan Turing Institute under the EPSRC grant EP/N510129/1.  ... 
doi:10.48550/arxiv.2111.00254 fatcat:4oaoly3cjbgpfe7dbmhqarwmmq