Home

Awesome

<!--lint ignore double-link-->

Awesome JAX Awesome<img src="https://raw.githubusercontent.com/google/jax/master/images/jax_logo_250px.png" alt="JAX Logo" align="right" height="100">

<!--lint ignore double-link-->

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.

<!--lint enable double-link-->

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

Contents

<a name="libraries" />

Libraries

<a name="new-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.

<a name="models-and-projects" />

Models and Projects

JAX

Flax

Haiku

Trax

NumPyro

<a name="videos" />

Videos

<a name="papers" />

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.

<!--lint ignore awesome-list-item--> <!--lint enable awesome-list-item--> <a name="tutorials-and-blog-posts" />

Tutorials and Blog Posts

<a name="books" />

Books

<a name="community" />

Community

Contributing

Contributions welcome! Read the contribution guidelines first.