Track Awesome Jax Updates Weekly

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

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

[ Daily / Weekly / Overview ]

Jan 05 - Jan 11, 2026

Libraries

• Dynamax (⭐902) - Probabilistic state space models.

Libraries / New Libraries

• astronomix (⭐44) - differentiable (magneto)hydrodynamics for astrophysics in JAX

• vivsim (⭐28) - Fluid-structure interaction simulations using Immersed Boundary-Lattice Boltzmann Method.

• MBIRJAX (⭐19) - High-performance tomographic reconstruction.

Sep 08 - Sep 14, 2025

Libraries / New Libraries

• flaxdiff (⭐40) - Framework and Library for building and training Diffusion models in multi-node multi-device distributed settings (TPUs)

• FDTDX (⭐191) - Finite-Difference Time-Domain Electromagnetic Simulations in JAX

• DiffeRT (⭐46) - Differentiable Ray Tracing toolbox for Radio Propagation powered by the JAX ecosystem.

• JAX-in-Cell (⭐22) - Plasma physics simulations using a PIC (Particle-in-Cell) method to self-consistently solve for electron and ion dynamics in electromagnetic fields

• kvax (⭐153) - A FlashAttention implementation for JAX with support for efficient document mask computation and context parallelism.

Models and Projects / Flax

• awesome-jax-flax-llms - Collection of LLMs implemented in JAX & Flax

Models and Projects / Equinox

• Sampling Path Candidates with Machine Learning - Official tutorial and implementation from the paper Towards Generative Ray Path Sampling for Faster Point-to-Point Ray Tracing.

Feb 17 - Feb 23, 2025

Libraries / New Libraries

• foragax (⭐5) - Agent-Based modelling framework in JAX.

Models and Projects / Flax

• DeepSeek-R1-Flax-1.5B-Distill (⭐26) - Flax implementation of DeepSeek-R1 1.5B distilled reasoning LLM.

• JaxNeuS (⭐1) - Implementation of NeuS: Learning Neural Implicit Surfaces by Volume Rendering for Multi-view Reconstruction

• DETR (⭐8) - Flax implementation of DETR: End-to-end Object Detection with Transformers using Sinkhorn solver and parallel bipartite matching.

Feb 10 - Feb 16, 2025

Libraries / New Libraries

• tmmax (⭐26) - Vectorized calculation of optical properties in thin-film structures using JAX. Swiss Army knife tool for thin-film optics research

Dec 02 - Dec 08, 2024

Libraries / New Libraries

• Brain Dynamics Programming Ecosystem
  • BrainPy (⭐645) - Brain Dynamics Programming in Python.
  • brainunit (⭐15) - Physical units and unit-aware mathematical system in JAX.
  • dendritex (⭐10) - Dendritic Modeling in JAX.
  • brainstate (⭐19) - State-based Transformation System for Program Compilation and Augmentation.
  • braintaichi (⭐6) - Leveraging Taichi Lang to customize brain dynamics operators.

Oct 14 - Oct 20, 2024

Tutorials and Blog Posts / Equinox

• Simple PDE solver + Constrained Optimization with JAX by Philip Mocz - A simple example of solving the advection-diffusion equations with JAX and using it in a constrained optimization problem to find initial conditions that yield desired result.

Sep 02 - Sep 08, 2024

Libraries / New Libraries

• Coreax (⭐36) - Algorithms for finding coresets to compress large datasets while retaining their statistical properties.

Aug 26 - Sep 01, 2024

Libraries / New Libraries

• Nonlinear Optimization
  • Optimistix (⭐522) - Root finding, minimisation, fixed points, and least squares.
  • JAXopt (⭐1k) - Hardware accelerated (GPU/TPU), batchable and differentiable optimizers in JAX.

Jun 17 - Jun 23, 2024

Libraries / New Libraries

• NAVIX (⭐154) - A reimplementation of MiniGrid, a Reinforcement Learning environment, in JAX

Apr 29 - May 05, 2024

Libraries

• Neural Network Libraries
  • Flax (⭐7k) - Centered on flexibility and clarity.
  • Flax NNX (⭐7k) - An evolution on Flax by the same team
  • Haiku (⭐3.2k) - Focused on simplicity, created by the authors of Sonnet at DeepMind.
  • Objax (⭐775) - Has an object oriented design similar to PyTorch.
  • Elegy - A High Level API for Deep Learning in JAX. Supports Flax, Haiku, and Optax.
  • Trax (⭐8.3k) - "Batteries included" deep learning library focused on providing solutions for common workloads.
  • Jraph (⭐1.5k) - Lightweight graph neural network library.
  • Neural Tangents (⭐2.4k) - High-level API for specifying neural networks of both finite and infinite width.
  • HuggingFace Transformers (⭐154k) - Ecosystem of pretrained Transformers for a wide range of natural language tasks (Flax).
  • Equinox (⭐2.7k) - Callable PyTrees and filtered JIT/grad transformations => neural networks in JAX.
  • Scenic (⭐3.7k) - A Jax Library for Computer Vision Research and Beyond.
  • Penzai (⭐1.8k) - Prioritizes legibility, visualization, and easy editing of neural network models with composable tools and a simple mental model.

Libraries / New Libraries

• flowjax (⭐205) - Distributions and normalizing flows built as equinox modules.

• EasyDeL (⭐328) - EasyDeL 🔮 is an OpenSource Library to make your training faster and more Optimized With cool Options for training and serving (Llama, MPT, Mixtral, Falcon, etc) in JAX

Apr 22 - Apr 28, 2024

Libraries / New Libraries

• dynamiqs (⭐261) - High-performance and differentiable simulations of quantum systems with JAX.

Mar 18 - Mar 24, 2024

Libraries

• Levanter (⭐687) - Legible, Scalable, Reproducible Foundation Models with Named Tensors and JAX.

• EasyLM (⭐2.5k) - LLMs made easy: Pre-training, finetuning, evaluating and serving LLMs in JAX/Flax.

Community / Equinox

• JaxLLM (Unofficial) Discord

Dec 18 - Dec 24, 2023

Libraries / New Libraries

• XLB (⭐427) - A Differentiable Massively Parallel Lattice Boltzmann Library in Python for Physics-Based Machine Learning.

Papers / Equinox

• XLB: A Differentiable Massively Parallel Lattice Boltzmann Library in Python. Mohammadmehdi Ataei, Hesam Salehipour. arXiv 2023. - White paper describing the XLB library: benchmarks, validations, and more details about the library.

Oct 23 - Oct 29, 2023

Libraries / New Libraries

• Spyx (⭐129) - Spiking Neural Networks in JAX for machine learning on neuromorphic hardware.

Oct 02 - Oct 08, 2023

Libraries / New Libraries

• Pgx (⭐566) - Vectorized board game environments for RL with an AlphaZero example.

Sep 25 - Oct 01, 2023

Libraries / New Libraries

• JAX Toolbox (⭐370) - Nightly CI and optimized examples for JAX on NVIDIA GPUs using libraries such as T5x, Paxml, and Transformer Engine.

Tutorials and Blog Posts / Equinox

• Learn JAX: From Linear Regression to Neural Networks by Rito Ghosh - A gentle introduction to JAX and using it to implement Linear and Logistic Regression, and Neural Network models and using them to solve real world problems.

Sep 18 - Sep 24, 2023

Libraries

• BlackJAX (⭐991) - Library of samplers for JAX.

Libraries / New Libraries

• QDax (⭐334) - Quality Diversity optimization in Jax.

Sep 11 - Sep 17, 2023

Libraries / New Libraries

• Praxis (⭐191) - The layer library for Pax with a goal to be usable by other JAX-based ML projects.

Aug 21 - Aug 27, 2023

Libraries / New Libraries

• Kernex (⭐71) - Differentiable stencil decorators in JAX.

• MaxText (⭐2.1k) - A simple, performant and scalable Jax LLM written in pure Python/Jax and targeting Google Cloud TPUs.

• Pax (⭐543) - A Jax-based machine learning framework for training large scale models.

Jul 31 - Aug 06, 2023

Libraries / New Libraries

• SCICO (⭐146) - Scientific computational imaging in JAX.

May 01 - May 07, 2023

Libraries / New Libraries

• purejaxrl (⭐1k) - Vectorisable, end-to-end RL algorithms in JAX.

• Lorax (⭐143) - Automatically apply LoRA to JAX models (Flax, Haiku, etc.)

Tutorials and Blog Posts / Equinox

• Achieving 4000x Speedups with PureJaxRL - A blog post on how JAX can massively speedup RL training through vectorisation.

Apr 10 - Apr 16, 2023

Libraries / New Libraries

• OTT-JAX (⭐681) - Optimal transport tools in JAX.

Feb 27 - Mar 05, 2023

Libraries / New Libraries

• safejax (⭐47) - Serialize JAX, Flax, Haiku, or Objax model params with 🤗safetensors.

Jan 16 - Jan 22, 2023

Libraries / New Libraries

• jax-tqdm (⭐124) - Add a tqdm progress bar to JAX scans and loops.

Jan 09 - Jan 15, 2023

Models and Projects / Flax

• GNNs for Solving Combinatorial Optimization Problems (⭐64) - A JAX + Flax implementation of Combinatorial Optimization with Physics-Inspired Graph Neural Networks.

Dec 26 - Jan 01, 2022

Libraries / New Libraries

• SPU (⭐314) - A domain-specific compiler and runtime suite to run JAX code with MPC(Secure Multi-Party Computation).

Dec 19 - Dec 25, 2022

Libraries

• Fortuna (⭐924) - AWS library for Uncertainty Quantification in Deep Learning.

Nov 14 - Nov 20, 2022

Libraries / New Libraries

• econpizza (⭐107) - Solve macroeconomic models with hetereogeneous agents using JAX.

Oct 17 - Oct 23, 2022

Models and Projects / NumPyro

• lqg (⭐30) - Official implementation of Bayesian inverse optimal control for linear-quadratic Gaussian problems from the paper Putting perception into action with inverse optimal control for continuous psychophysics

Sep 05 - Sep 11, 2022

Libraries / New Libraries

• JAXFit (⭐59) - Accelerated curve fitting library for nonlinear least-squares problems (see arXiv paper).

Aug 29 - Sep 04, 2022

Libraries / New Libraries

• GPJax (⭐566) - Gaussian processes in JAX.

• Jumanji (⭐790) - A Suite of Industry-Driven Hardware-Accelerated RL Environments written in JAX.

• Eqxvision (⭐109) - Equinox version of Torchvision.

Models and Projects / Flax

• Sharpened Cosine Similarity in JAX by Raphael Pisoni - A JAX/Flax implementation of the Sharpened Cosine Similarity layer.

Tutorials and Blog Posts / Equinox

• Writing a Training Loop in JAX + FLAX by Saurav Maheshkar and Soumik Rakshit - A tutorial on writing a simple end-to-end training and evaluation pipeline in JAX, Flax and Optax.

• Implementing NeRF in JAX by Soumik Rakshit and Saurav Maheshkar - A tutorial on 3D volumetric rendering of scenes represented by Neural Radiance Fields in JAX.

• 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.

Aug 22 - Aug 28, 2022

Books / Equinox

• Jax in Action - A hands-on guide to using JAX for deep learning and other mathematically-intensive applications.

Aug 08 - Aug 14, 2022

Libraries / New Libraries

• jwave (⭐186) - A library for differentiable acoustic simulations

Jun 27 - Jul 03, 2022

Libraries / New Libraries

• Mctx (⭐2.6k) - Monte Carlo tree search algorithms in native JAX.

• KFAC-JAX (⭐302) - Second Order Optimization with Approximate Curvature for NNs.

• TF2JAX (⭐118) - Convert functions/graphs to JAX functions.

Jun 20 - Jun 26, 2022

Libraries / New Libraries

• gymnax (⭐844) - Reinforcement Learning Environments with the well-known gym API.

Mar 14 - Mar 20, 2022

Libraries / New Libraries

• Einshape (⭐107) - DSL-based reshaping library for JAX and other frameworks.

• ALX (⭐37k) - Open-source library for distributed matrix factorization using Alternating Least Squares, more info in ALX: Large Scale Matrix Factorization on TPUs.

• Diffrax (⭐1.9k) - Numerical differential equation solvers in JAX.

• tinygp (⭐328) - The tiniest of Gaussian process libraries in JAX.

Models and Projects / JAX

• Symbolic Functionals (⭐37k) - Demonstration from Evolving symbolic density functionals.

• TriMap (⭐37k) - Official JAX implementation of TriMap: Large-scale Dimensionality Reduction Using Triplets.

Models and Projects / Flax

• mip-NeRF (⭐935) - Official implementation of Mip-NeRF: A Multiscale Representation for Anti-Aliasing Neural Radiance Fields.

• RegNeRF (⭐37k) - Official implementation of RegNeRF: Regularizing Neural Radiance Fields for View Synthesis from Sparse Inputs.

• ARDM (⭐37k) - Official implementation of Autoregressive Diffusion Models.

• D3PM (⭐37k) - Official implementation of Structured Denoising Diffusion Models in Discrete State-Spaces.

• Gumbel-max Causal Mechanisms (⭐37k) - Code for Learning Generalized Gumbel-max Causal Mechanisms, with extra code in GuyLor/gumbel_max_causal_gadgets_part2 (⭐2).

• Latent Programmer (⭐37k) - Code for the ICML 2021 paper Latent Programmer: Discrete Latent Codes for Program Synthesis.

• SNeRG (⭐37k) - Official implementation of Baking Neural Radiance Fields for Real-Time View Synthesis.

• Spin-weighted Spherical CNNs (⭐37k) - Adaptation of Spin-Weighted Spherical CNNs.

• VDVAE (⭐37k) - Adaptation of Very Deep VAEs Generalize Autoregressive Models and Can Outperform Them on Images, original code at openai/vdvae (⭐449).

• MUSIQ (⭐37k) - Checkpoints and model inference code for the ICCV 2021 paper MUSIQ: Multi-scale Image Quality Transformer

• AQuaDem (⭐37k) - Official implementation of Continuous Control with Action Quantization from Demonstrations.

• Combiner (⭐37k) - Official implementation of Combiner: Full Attention Transformer with Sparse Computation Cost.

• Dreamfields (⭐37k) - Official implementation of the ICLR 2022 paper Progressive Distillation for Fast Sampling of Diffusion Models.

• GIFT (⭐37k) - Official implementation of Gradual Domain Adaptation in the Wild:When Intermediate Distributions are Absent.

• Light Field Neural Rendering (⭐37k) - Official implementation of Light Field Neural Rendering.

Feb 28 - Mar 06, 2022

Libraries / New Libraries

• Neural Network Libraries
  • FedJAX (⭐267) - Federated learning in JAX, built on Optax and Haiku.
  • Equivariant MLP (⭐281) - Construct equivariant neural network layers.
  • jax-resnet (⭐118) - Implementations and checkpoints for ResNet variants in Flax.
  • jax-raft (⭐3) - JAX/Flax port of the RAFT optical flow estimator.
  • Parallax (⭐153) - Immutable Torch Modules for JAX.

Feb 21 - Feb 27, 2022

Libraries / New Libraries

• evosax (⭐712) - JAX-Based Evolution Strategies

• SymJAX (⭐125) - Symbolic CPU/GPU/TPU programming.

• mcx (⭐329) - Express & compile probabilistic programs for performant inference.

Feb 14 - Feb 20, 2022

Libraries / New Libraries

• EvoJAX (⭐929) - Hardware-Accelerated Neuroevolution

Feb 07 - Feb 13, 2022

Libraries / New Libraries

• PGMax (⭐65) - A framework for building discrete Probabilistic Graphical Models (PGM's) and running inference inference on them via JAX.

Tutorials and Blog Posts / Equinox

• Get started with JAX by Aleksa Gordić (⭐776) - A series of notebooks and videos going from zero JAX knowledge to building neural networks in Haiku.

Jan 24 - Jan 30, 2022

Libraries / New Libraries

• tree-math (⭐205) - Convert functions that operate on arrays into functions that operate on PyTrees.

• jax-models (⭐160) - Implementations of research papers originally without code or code written with frameworks other than JAX.

Jan 03 - Jan 09, 2022

Tutorials and Blog Posts / Equinox

• Introduction to JAX by Kevin Murphy - Colab that introduces various aspects of the language and applies them to simple ML problems.

Dec 27 - Jan 02, 2021

Models and Projects / Haiku

• Two Player Auction Learning (⭐0) - JAX implementation of the paper Auction learning as a two-player game.

Nov 08 - Nov 14, 2021

Libraries / New Libraries

• JaxDF (⭐131) - Framework for differentiable simulators with arbitrary discretizations.

Sep 06 - Sep 12, 2021

Libraries / New Libraries

• bayex (⭐100) - Bayesian Optimization powered by JAX.

Jul 26 - Aug 01, 2021

Models and Projects / JAX

• Amortized Bayesian Optimization (⭐37k) - Code related to Amortized Bayesian Optimization over Discrete Spaces.

• Accurate Quantized Training (⭐37k) - Tools and libraries for running and analyzing neural network quantization experiments in JAX and Flax.

• BNN-HMC (⭐37k) - Implementation for the paper What Are Bayesian Neural Network Posteriors Really Like?.

• JAX-DFT (⭐37k) - One-dimensional density functional theory (DFT) in JAX, with implementation of Kohn-Sham equations as regularizer: building prior knowledge into machine-learned physics.

• Robust Loss (⭐37k) - Reference code for the paper A General and Adaptive Robust Loss Function.

Models and Projects / Flax

• Performer (⭐37k) - Flax implementation of the Performer (linear transformer via FAVOR+) architecture.

• JaxNeRF (⭐37k) - Implementation of NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis with multi-device GPU/TPU support.

• Big Transfer (BiT) (⭐1.5k) - Implementation of Big Transfer (BiT): General Visual Representation Learning.

• JAX RL (⭐740) - Implementations of reinforcement learning algorithms.

• gMLP - Implementation of Pay Attention to MLPs.

• MLP Mixer - Minimal implementation of MLP-Mixer: An all-MLP Architecture for Vision.

• Distributed Shampoo (⭐37k) - Implementation of Second Order Optimization Made Practical.

• NesT (⭐201) - Official implementation of Aggregating Nested Transformers.

• XMC-GAN (⭐96) - Official implementation of Cross-Modal Contrastive Learning for Text-to-Image Generation.

• FNet (⭐37k) - Official implementation of FNet: Mixing Tokens with Fourier Transforms.

• GFSA (⭐37k) - Official implementation of Learning Graph Structure With A Finite-State Automaton Layer.

• IPA-GNN (⭐37k) - Official implementation of Learning to Execute Programs with Instruction Pointer Attention Graph Neural Networks.

• Flax Models (⭐37k) - Collection of models and methods implemented in Flax.

• Protein LM (⭐37k) - Implements BERT and autoregressive models for proteins, as described in Biological Structure and Function Emerge from Scaling Unsupervised Learning to 250 Million Protein Sequences and ProGen: Language Modeling for Protein Generation.

• Slot Attention (⭐37k) - Reference implementation for Differentiable Patch Selection for Image Recognition.

• Vision Transformer (⭐12k) - Official implementation of An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale.

• FID computation (⭐29) - Port of mseitzer/pytorch-fid (⭐3.8k) to Flax.

Models and Projects / Haiku

• AlphaFold (⭐14k) - Implementation of the inference pipeline of AlphaFold v2.0, presented in Highly accurate protein structure prediction with AlphaFold.

• Adversarial Robustness (⭐15k) - Reference code for Uncovering the Limits of Adversarial Training against Norm-Bounded Adversarial Examples and Fixing Data Augmentation to Improve Adversarial Robustness.

• Bootstrap Your Own Latent (⭐15k) - Implementation for the paper Bootstrap your own latent: A new approach to self-supervised Learning.

• Gated Linear Networks (⭐15k) - GLNs are a family of backpropagation-free neural networks.

• Glassy Dynamics (⭐15k) - Open source implementation of the paper Unveiling the predictive power of static structure in glassy systems.

• MMV (⭐15k) - Code for the models in Self-Supervised MultiModal Versatile Networks.

• Normalizer-Free Networks (⭐15k) - Official Haiku implementation of NFNets.

• NuX (⭐86) - Normalizing flows with JAX.

• OGB-LSC (⭐15k) - This repository contains DeepMind's entry to the PCQM4M-LSC (quantum chemistry) and MAG240M-LSC (academic graph) tracks of the OGB Large-Scale Challenge (OGB-LSC).

• Persistent Evolution Strategies (⭐37k) - Code used for the paper Unbiased Gradient Estimation in Unrolled Computation Graphs with Persistent Evolution Strategies.

• WikiGraphs (⭐15k) - Baseline code to reproduce results in WikiGraphs: A Wikipedia Text - Knowledge Graph Paired Datase.

Models and Projects / Trax

• Reformer (⭐8.3k) - Implementation of the Reformer (efficient transformer) architecture.

Videos / Equinox

• JAX, Flax & Transformers 🤗 (⭐154k) - 3 days of talks around JAX / Flax, Transformers, large-scale language modeling and other great topics.

Jul 19 - Jul 25, 2021

Libraries / New Libraries

• PIX (⭐431) - PIX is an image processing library in JAX, for JAX.

Jul 12 - Jul 18, 2021

Libraries

• NetKet (⭐649) - Machine Learning toolbox for Quantum Physics.

Jun 07 - Jun 13, 2021

Libraries / New Libraries

• flaxmodels (⭐263) - Pretrained models for Jax/Flax.

• CR.Sparse (⭐96) - XLA accelerated algorithms for sparse representations and compressive sensing.

• exojax (⭐66) - Automatic differentiable spectrum modeling of exoplanets/brown dwarfs compatible to JAX.

Tutorials and Blog Posts / Equinox

• Writing an MCMC sampler in JAX by Jeremie Coullon - Tutorial on the different ways to write an MCMC sampler in JAX along with speed benchmarks.

• How to add a progress bar to JAX scans and loops by Jeremie Coullon - Tutorial on how to add a progress bar to compiled loops in JAX using the host_callback module.

May 31 - Jun 06, 2021

Libraries / New Libraries

• delta PV (⭐64) - A photovoltaic simulator with automatic differentation.

• jaxlie (⭐315) - Lie theory library for rigid body transformations and optimization.

• BRAX (⭐3k) - Differentiable physics engine to simulate environments along with learning algorithms to train agents for these environments.

May 17 - May 23, 2021

Libraries

• Coax (⭐181) - Turn RL papers into code, the easy way.

Apr 26 - May 02, 2021

Libraries

• Distrax (⭐613) - Reimplementation of TensorFlow Probability, containing probability distributions and bijectors.

• cvxpylayers (⭐2k) - Construct differentiable convex optimization layers.

• TensorLy (⭐1.7k) - Tensor learning made simple.

Mar 22 - Mar 28, 2021

Libraries / New Libraries

• Optimal Transport Tools (⭐214) - Toolbox that bundles utilities to solve optimal transport problems.

Mar 01 - Mar 07, 2021

Libraries

• RLax (⭐1.4k) - Library for implementing reinforcement learning agents.

Libraries / New Libraries

• FlaxVision (⭐45) - Flax version of TorchVision.

• Oryx (⭐4.4k) - Probabilistic programming language based on program transformations.

Videos / Equinox

• Solving y=mx+b with Jax on a TPU Pod slice - Mat Kelcey - A four part YouTube tutorial series with Colab notebooks that starts with Jax fundamentals and moves up to training with a data parallel approach on a v3-32 TPU Pod slice.

Tutorials and Blog Posts / Equinox

• Deterministic ADVI in JAX by Martin Ingram - Walk through of implementing automatic differentiation variational inference (ADVI) easily and cleanly with JAX.

• 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.

Feb 22 - Feb 28, 2021

Libraries / New Libraries

• imax (⭐41) - Image augmentations and transformations.

Feb 08 - Feb 14, 2021

Tutorials and Blog Posts / Equinox

• Evolving Neural Networks in JAX by Robert Tjarko Lange - Explores how JAX can power the next generation of scalable neuroevolution algorithms.

• Exploring hyperparameter meta-loss landscapes with JAX by Luke Metz - Demonstrates how to use JAX to perform inner-loss optimization with SGD and Momentum, outer-loss optimization with gradients, and outer-loss optimization using evolutionary strategies.

Feb 01 - Feb 07, 2021

Tutorials and Blog Posts / Equinox

• Out of distribution (OOD) detection by Mat Kelcey - Implements different methods for OOD detection.

• Extending JAX with custom C++ and CUDA code by Dan Foreman-Mackey (⭐402) - Tutorial demonstrating the infrastructure required to provide custom ops in JAX.

Jan 25 - Jan 31, 2021

Libraries

• NumPyro (⭐2.6k) - Probabilistic programming based on the Pyro library.

• Chex (⭐912) - Utilities to write and test reliable JAX code.

• Optax (⭐2.1k) - Gradient processing and optimization library.

• JAX, M.D. (⭐1.4k) - Accelerated, differential molecular dynamics.

Libraries / New Libraries

• jax-unirep (⭐108) - Library implementing the UniRep model for protein machine learning applications.

• jax-cosmo (⭐221) - Differentiable cosmology library.

Videos / Equinox

• Introduction to JAX - Simple neural network from scratch in JAX.

Papers / Equinox

• 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.

Tutorials and Blog Posts / Equinox

• Plugging Into JAX by Nick Doiron - Compares Flax, Haiku, and Objax on the Kaggle flower classification challenge.

Jan 18 - Jan 24, 2021

Models and Projects / JAX

• Fourier Feature Networks (⭐1.4k) - Official implementation of Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains.

Videos / Equinox

• NeurIPS 2020: JAX Ecosystem Meetup - JAX, its use at DeepMind, and discussion between engineers, scientists, and JAX core team.

• JAX: Accelerated Machine Learning Research | SciPy 2020 | VanderPlas - JAX's core design, how it's powering new research, and how you can start using it.

• Bayesian Programming with JAX + NumPyro — Andy Kitchen - Introduction to Bayesian modelling using NumPyro.

• JAX on Cloud TPUs | NeurIPS 2020 | Skye Wanderman-Milne and James Bradbury - Presentation of TPU host access with demo.

• Deep Implicit Layers - Neural ODEs, Deep Equilibirum Models, and Beyond | NeurIPS 2020 - Tutorial created by Zico Kolter, David Duvenaud, and Matt Johnson with Colab notebooks avaliable in Deep Implicit Layers.

Papers / Equinox

• 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.

• 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.

Tutorials and Blog Posts / Equinox

• Using JAX to accelerate our research by David Budden and Matteo Hessel - Describes the state of JAX and the JAX ecosystem at DeepMind.

• Getting started with JAX (MLPs, CNNs & RNNs) by Robert Lange - Neural network building blocks from scratch with the basic JAX operators.

• Tutorial: image classification with JAX and Flax Linen by 8bitmp3 (⭐24) - Learn how to create a simple convolutional network with the Linen API by Flax and train it to recognize handwritten digits.

• Meta-Learning in 50 Lines of JAX by Eric Jang - Introduction to both JAX and Meta-Learning.

• Normalizing Flows in 100 Lines of JAX by Eric Jang - Concise implementation of RealNVP.

• 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.

• Understanding Autodiff with JAX by Srihari Radhakrishna - Understand how autodiff works using JAX.

• From PyTorch to JAX: towards neural net frameworks that purify stateful code by Sabrina J. Mielke - Showcases how to go from a PyTorch-like style of coding to a more Functional-style of coding.

Jan 04 - Jan 10, 2021

Libraries / New Libraries

• mpi4jax (⭐507) - Combine MPI operations with your Jax code on CPUs and GPUs.

Jan 04 - Jan 10, 2021

Libraries / New Libraries

• jax-flows - Normalizing flows in JAX.

• sklearn-jax-kernels (⭐45) - scikit-learn kernel matrices using JAX.

• efax (⭐75) - Exponential Families in JAX.

Models and Projects / JAX

• kalman-jax (⭐102) - Approximate inference for Markov (i.e., temporal) Gaussian processes using iterated Kalman filtering and smoothing.

• jaxns (⭐218) - Nested sampling in JAX.

Dec 28 - Jan 03, 2020

Videos / Equinox

• JAX: Accelerated machine-learning research via composable function transformations in Python | NeurIPS 2019 | Skye Wanderman-Milne - JAX intro presentation in Program Transformations for Machine Learning workshop.

Community / Equinox

• JAX GitHub Discussions (⭐34k)

• Reddit