JAX - A curated list of resources https://github.com/google/jax
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!
This section contains libraries that are well-made and useful, but have not necessarily been battle-tested by a large userbase yet.
scikit-learnkernel matrices using JAX.
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.
Contributions welcome! Read the contribution guidelines first.