From PyTorch to JAX: towards neural net frameworks that purify stateful code by Sabrina J. Mielke

sjmielke.comsite

Showcases how to go from a PyTorch-like style of coding to a more Functional-style of coding.

from

JAX

added

2026-10-10

likes

0

similar

1. ▲

   Evolving Neural Networks in JAX by Robert Tjarko Lange roberttlange.github.io

   Explores how JAX can power the next generation of scalable neuroevolution algorithms.

   site topic aitopic programming link ↗
   site
2. ▲

   Objax github.com

   Has an object oriented design similar to PyTorch.

   tool topic aitopic design link ↗
   
3. ▲

   Equinox github.com

   Callable PyTrees and filtered JIT/grad transformations => neural networks in JAX.

   tool topic ai link ↗
   
4. ▲

   Get started with JAX by Aleksa Gordić github.com

   A series of notebooks and videos going from zero JAX knowledge to building neural networks in Haiku.

   tool topic ai link ↗
   
5. ▲

   JAX: Accelerated machine-learning research via composable function transformations in Python | NeurIPS 2019 | Skye… slideslive.com

   JAX intro presentation in Program Transformations for Machine Learning workshop.

   site topic aitopic programming link ↗
   
6. ▲

   TF2JAX github.com

   Convert functions/graphs to JAX functions.

   tool link ↗
   

JAX › Tutorials and Blog Posts: “Showcases how to go from a PyTorch-like style of coding to a more Functional-style of coding.”