Awesome Jax Overview

JAX - A curated list of resources https://github.com/google/jax

🏠 Home · 🔥 Feed · 📮 Subscribe · ❤️ Sponsor · 😺 n2cholas/awesome-jax · ⭐ 1.3K · 🏷️ Computer Science

[ Daily / Weekly / Overview ]

Awesome JAX AwesomeJAX Logo

JAX (⭐28k) 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

NumPyro

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

Books

Community

Contributing

Contributions welcome! Read the contribution guidelines first.