Awesome JAX Awesome

JAX brings automatic differentiation and the XLA compiler together through a NumPy-like API for high performance machine learning research on accelerators like GPUs and TPUs.

This is a curated list of awesome JAX libraries, projects, and other resources. Contributions are welcome!

Contents

Libraries

New Libraries

This section contains libraries that are well-made and useful, but have not necessarily been battle-tested by a large userbase yet.

Models and Projects

JAX

Flax

Haiku

Trax

Videos

Papers

This section contains papers focused on JAX (e.g. JAX-based library whitepapers, research on JAX, etc). Papers implemented in JAX are listed in the Models/Projects section.

Tutorials and Blog Posts

Community

Contributing

Contributions welcome! Read the contribution guidelines first.