awesome-jax
github.com/n2cholas/awesome-jax ↗JAX - A curated list of resources https://github.com/google/jax
Use this list with your AI agent
Add the Context Awesome MCP server to Claude, Cursor, or any MCP client, then ask:
"Show me jax resources from awesome-jax"
Installation instructions →What's inside
Models and Projects
- Accurate Quantized TrainingJAX
Tools and libraries for running and analyzing neural network quantization experiments in JAX and Flax.
- Adversarial RobustnessHaiku
Reference code for
- AlphaFoldHaiku
Implementation of the inference pipeline of AlphaFold v2.0, presented in
- Amortized Bayesian OptimizationJAX
Code related to
- AQuaDemFlax
Official implementation of
- ARDMFlax
Official implementation of
Tutorials and Blog Posts
- Achieving 4000x Speedups with PureJaxRL
A blog post on how JAX can massively speedup RL training through vectorisation.
- Deep Learning tutorials with JAX+Flax by Phillip Lippe
A series of notebooks explaining various deep learning concepts, from basics (e.g. intro to JAX/Flax, activiation functions) to recent advances (e.g., Vision Transformers, SimCLR), with translations to PyTorch.
- Deterministic ADVI in JAX by Martin Ingram
Walk through of implementing automatic differentiation variational inference (ADVI) easily and cleanly with JAX.
- Differentiable Path Tracing on the GPU/TPU by Eric Jang
Tutorial on implementing path tracing.
- Ensemble networks by Mat Kelcey
Ensemble nets are a method of representing an ensemble of models as one single logical model.
- Evolved channel selection by Mat Kelcey
Trains a classification model robust to different combinations of input channels at different resolutions, then uses a genetic algorithm to decide the best combination for a particular loss.
Libraries
- ALXNew Libraries
Open-source library for distributed matrix factorization using Alternating Least Squares, more info in
- astronomixNew Libraries
differentiable (magneto)hydrodynamics for astrophysics in JAX
- bayexNew Libraries
Bayesian Optimization powered by JAX.
- BlackJAX
Library of samplers for JAX.
- BrainPyNew Libraries
Brain Dynamics Programming in Python.
- brainstateNew Libraries
State-based Transformation System for Program Compilation and Augmentation.
Videos
- Bayesian Programming with JAX + NumPyro — Andy Kitchen
Introduction to Bayesian modelling using NumPyro.
- Deep Implicit Layers - Neural ODEs, Deep Equilibirum Models, and Beyond | NeurIPS 2020
NeurIPS 2020 - Tutorial created by Zico Kolter, David Duvenaud, and Matt Johnson with Colab notebooks avaliable in Deep Implicit Layers.
- Introduction to JAX
Simple neural network from scratch in JAX.
- JAX: Accelerated Machine Learning Research | SciPy 2020 | VanderPlas
SciPy 2020 | VanderPlas - JAX's core design, how it's powering new research, and how you can start using it.
- JAX: Accelerated machine-learning research via composable function transformations in Python | NeurIPS 2019 | Skye Wanderman-Milne
NeurIPS 2019 | Skye Wanderman-Milne - JAX intro presentation in Program Transformations for Machine Learning workshop.
- JAX, Flax & Transformers 🤗
3 days of talks around JAX / Flax, Transformers, large-scale language modeling and other great topics.
Papers
- Compiling machine learning programs via high-level tracing. Roy Frostig, Matthew James Johnson, Chris Leary. MLSys 2018.
White paper describing an early version of JAX, detailing how computation is traced and compiled.
- Enabling Fast Differentially Private SGD via Just-in-Time Compilation and Vectorization. Pranav Subramani, Nicholas Vadivelu, Gautam Kamath. arXiv 2020.
Uses JAX's JIT and VMAP to achieve faster differentially private than existing libraries.
- JAX, M.D.: A Framework for Differentiable Physics. Samuel S. Schoenholz, Ekin D. Cubuk. NeurIPS 2020.
Introduces JAX, M.D., a differentiable physics library which includes simulation environments, interaction potentials, neural networks, and more.
Books
- Jax in Action
A hands-on guide to using JAX for deep learning and other mathematically-intensive applications.
Showing a sample of 217 resources. View the full list on GitHub →