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-learn
kernel matrices using JAX.
safetensors
.
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.
host_callback
module.Contributions welcome! Read the contribution guidelines first.