Community

Equinox

• Reddit
• JaxLLM (Unofficial) Discord
• JAX GitHub Discussions

Libraries

New Libraries

• ALX - Open-source library for distributed matrix factorization using Alternating Least Squares, more info in ALX: Large Scale Matrix Factorization on TPUs.
• dynamiqs - High-performance and differentiable simulations of quantum systems with JAX.
• FedJAX - Federated learning in JAX, built on Optax and Haiku.
• Equivariant MLP - Construct equivariant neural network layers.
• Parallax - Immutable Torch Modules for JAX.
• jax-unirep - Library implementing the UniRep model for protein machine learning applications.
• jax-flows - Normalizing flows in JAX.
• sklearn-jax-kernels - scikit-learn kernel matrices using JAX.
• jax-cosmo - Differentiable cosmology library.
• efax - Exponential Families in JAX.
• imax - Image augmentations and transformations.
• FlaxVision - Flax version of TorchVision.
• Optimal Transport Tools - Toolbox that bundles utilities to solve optimal transport problems.
• delta PV - A photovoltaic simulator with automatic differentation.
• jaxlie - Lie theory library for rigid body transformations and optimization.
• BRAX - Differentiable physics engine to simulate environments along with learning algorithms to train agents for these environments.
• flaxmodels - Pretrained models for Jax/Flax.
• CR.Sparse - XLA accelerated algorithms for sparse representations and compressive sensing.
• exojax - Automatic differentiable spectrum modeling of exoplanets/brown dwarfs compatible to JAX.
• JAXopt - Hardware accelerated (GPU/TPU), batchable and differentiable optimizers in JAX.
• bayex - Bayesian Optimization powered by JAX.
• JaxDF - Framework for differentiable simulators with arbitrary discretizations.
• tree-math - Convert functions that operate on arrays into functions that operate on PyTrees.
• jax-models - Implementations of research papers originally without code or code written with frameworks other than JAX.
• PGMax - A framework for building discrete Probabilistic Graphical Models (PGM's) and running inference inference on them via JAX.
• EvoJAX - Hardware-Accelerated Neuroevolution
• evosax - JAX-Based Evolution Strategies
• SymJAX - Symbolic CPU/GPU/TPU programming.
• mcx - Express & compile probabilistic programs for performant inference.
• Diffrax - Numerical differential equation solvers in JAX.
• tinygp - The tiniest of Gaussian process libraries in JAX.
• gymnax - Reinforcement Learning Environments with the well-known gym API.
• jwave - A library for differentiable acoustic simulations
• Jumanji - A Suite of Industry-Driven Hardware-Accelerated RL Environments written in JAX.
• Eqxvision - Equinox version of Torchvision.
• JAXFit - Accelerated curve fitting library for nonlinear least-squares problems (see arXiv paper).
• econpizza - Solve macroeconomic models with hetereogeneous agents using JAX.
• SPU - A domain-specific compiler and runtime suite to run JAX code with MPC(Secure Multi-Party Computation).
• jax-tqdm - Add a tqdm progress bar to JAX scans and loops.
• safejax - Serialize JAX, Flax, Haiku, or Objax model params with 🤗safetensors.
• Kernex - Differentiable stencil decorators in JAX.
• Pax - A Jax-based machine learning framework for training large scale models.
• Praxis - The layer library for Pax with a goal to be usable by other JAX-based ML projects.
• purejaxrl - Vectorisable, end-to-end RL algorithms in JAX.
• Lorax - Automatically apply LoRA to JAX models (Flax, Haiku, etc.)
• SCICO - Scientific computational imaging in JAX.
• Spyx - Spiking Neural Networks in JAX for machine learning on neuromorphic hardware.
• BrainPy - Brain Dynamics Programming in Python.
• OTT-JAX - Optimal transport tools in JAX.
• QDax - Quality Diversity optimization in Jax.
• JAX Toolbox - Nightly CI and optimized examples for JAX on NVIDIA GPUs using libraries such as T5x, Paxml, and Transformer Engine.
• Optimistix - Root finding, minimisation, fixed points, and least squares.
• flowjax - Distributions and normalizing flows built as equinox modules.
• brainunit - Physical units and unit-aware mathematical system in JAX.
• brainstate - State-based Transformation System for Program Compilation and Augmentation.
• braintaichi - Leveraging Taichi Lang to customize brain dynamics operators.
• jax-resnet - Implementations and checkpoints for ResNet variants in Flax.
• mpi4jax - Combine MPI operations with your Jax code on CPUs and GPUs.
• PIX - PIX is an image processing library in JAX, for JAX.
• dendritex - Dendritic Modeling in JAX.
• flaxdiff - Framework and Library for building and training Diffusion models in multi-node multi-device distributed settings (TPUs)
• TF2JAX - Convert functions/graphs to JAX functions.
• GPJax - Gaussian processes in JAX.
• MaxText - A simple, performant and scalable Jax LLM written in pure Python/Jax and targeting Google Cloud TPUs.
• EasyDeL - 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
• XLB - A Differentiable Massively Parallel Lattice Boltzmann Library in Python for Physics-Based Machine Learning.
• foragax - Agent-Based modelling framework in JAX.
• tmmax - Vectorized calculation of optical properties in thin-film structures using JAX. Swiss Army knife tool for thin-film optics research
• Coreax - Algorithms for finding coresets to compress large datasets while retaining their statistical properties.
• NAVIX - A reimplementation of MiniGrid, a Reinforcement Learning environment, in JAX
• FDTDX - Finite-Difference Time-Domain Electromagnetic Simulations in JAX
• DiffeRT - Differentiable Ray Tracing toolbox for Radio Propagation powered by the JAX ecosystem.
• JAX-in-Cell - Plasma physics simulations using a PIC (Particle-in-Cell) method to self-consistently solve for electron and ion dynamics in electromagnetic fields
• kvax - A FlashAttention implementation for JAX with support for efficient document mask computation and context parallelism.
• Einshape - DSL-based reshaping library for JAX and other frameworks.
• KFAC-JAX - Second Order Optimization with Approximate Curvature for NNs.
• Oryx - Probabilistic programming language based on program transformations.
• Mctx - Monte Carlo tree search algorithms in native JAX.
• Pgx - Vectorized board game environments for RL with an AlphaZero example.
• astronomix - differentiable (magneto)hydrodynamics for astrophysics in JAX
• vivsim - Fluid-structure interaction simulations using Immersed Boundary-Lattice Boltzmann Method.
• MBIRJAX - High-performance tomographic reconstruction.

• Flax - Centered on flexibility and clarity.
• Objax - Has an object oriented design similar to PyTorch.
• Elegy - A High Level API for Deep Learning in JAX. Supports Flax, Haiku, and Optax.
• Trax - "Batteries included" deep learning library focused on providing solutions for common workloads.
• Neural Tangents - High-level API for specifying neural networks of both finite and infinite width.
• HuggingFace Transformers - Ecosystem of pretrained Transformers for a wide range of natural language tasks (Flax).
• Equinox - Callable PyTrees and filtered JIT/grad transformations => neural networks in JAX.
• Scenic - A Jax Library for Computer Vision Research and Beyond.
• EasyLM - LLMs made easy: Pre-training, finetuning, evaluating and serving LLMs in JAX/Flax.
• NumPyro - Probabilistic programming based on the Pyro library.
• Coax - Turn RL papers into code, the easy way.
• cvxpylayers - Construct differentiable convex optimization layers.
• TensorLy - Tensor learning made simple.
• NetKet - Machine Learning toolbox for Quantum Physics.
• Fortuna - AWS library for Uncertainty Quantification in Deep Learning.
• BlackJAX - Library of samplers for JAX.
• Flax NNX - An evolution on Flax by the same team
• Penzai - Prioritizes legibility, visualization, and easy editing of neural network models with composable tools and a simple mental model.
• Jraph - Lightweight graph neural network library.
• JAX, M.D. - Accelerated, differential molecular dynamics.
• RLax - Library for implementing reinforcement learning agents.
• Distrax - Reimplementation of TensorFlow Probability, containing probability distributions and bijectors.
• Dynamax - Probabilistic state space models.
• Haiku - Focused on simplicity, created by the authors of Sonnet at DeepMind.

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.

Flax

• Performer - Flax implementation of the Performer (linear transformer via FAVOR+) architecture.
• JaxNeRF - Implementation of NeRF: Representing Scenes as Neural Radiance Fields for View Synthesis with multi-device GPU/TPU support.
• RegNeRF - Official implementation of RegNeRF: Regularizing Neural Radiance Fields for View Synthesis from Sparse Inputs.
• gMLP - Implementation of Pay Attention to MLPs.
• MLP Mixer - Minimal implementation of MLP-Mixer: An all-MLP Architecture for Vision.
• Distributed Shampoo - Implementation of Second Order Optimization Made Practical.
• FNet - Official implementation of FNet: Mixing Tokens with Fourier Transforms.
• GFSA - Official implementation of Learning Graph Structure With A Finite-State Automaton Layer.
• IPA-GNN - Official implementation of Learning to Execute Programs with Instruction Pointer Attention Graph Neural Networks.
• Flax Models - Collection of models and methods implemented in Flax.
• Protein LM - 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 - Reference implementation for Differentiable Patch Selection for Image Recognition.
• ARDM - Official implementation of Autoregressive Diffusion Models.
• D3PM - Official implementation of Structured Denoising Diffusion Models in Discrete State-Spaces.
• Gumbel-max Causal Mechanisms - Code for Learning Generalized Gumbel-max Causal Mechanisms, with extra code in GuyLor/gumbel_max_causal_gadgets_part2.
• Latent Programmer - Code for the ICML 2021 paper Latent Programmer: Discrete Latent Codes for Program Synthesis.
• SNeRG - Official implementation of Baking Neural Radiance Fields for Real-Time View Synthesis.
• Spin-weighted Spherical CNNs - Adaptation of Spin-Weighted Spherical CNNs.
• VDVAE - Adaptation of Very Deep VAEs Generalize Autoregressive Models and Can Outperform Them on Images, original code at openai/vdvae.
• MUSIQ - Checkpoints and model inference code for the ICCV 2021 paper MUSIQ: Multi-scale Image Quality Transformer
• AQuaDem - Official implementation of Continuous Control with Action Quantization from Demonstrations.
• Combiner - Official implementation of Combiner: Full Attention Transformer with Sparse Computation Cost.
• Dreamfields - Official implementation of the ICLR 2022 paper Progressive Distillation for Fast Sampling of Diffusion Models.
• GIFT - Official implementation of Gradual Domain Adaptation in the Wild:When Intermediate Distributions are Absent.
• Light Field Neural Rendering - Official implementation of Light Field Neural Rendering.
• Sharpened Cosine Similarity in JAX by Raphael Pisoni - A JAX/Flax implementation of the Sharpened Cosine Similarity layer.
• mip-NeRF - Official implementation of Mip-NeRF: A Multiscale Representation for Anti-Aliasing Neural Radiance Fields.
• JaxNeuS - Implementation of NeuS: Learning Neural Implicit Surfaces by Volume Rendering for Multi-view Reconstruction
• Big Transfer (BiT) - Implementation of Big Transfer (BiT): General Visual Representation Learning.
• NesT - Official implementation of Aggregating Nested Transformers.
• XMC-GAN - Official implementation of Cross-Modal Contrastive Learning for Text-to-Image Generation.
• Vision Transformer - Official implementation of An Image is Worth 16x16 Words: Transformers for Image Recognition at Scale.
• FID computation - Port of mseitzer/pytorch-fid to Flax.
• GNNs for Solving Combinatorial Optimization Problems - A JAX + Flax implementation of Combinatorial Optimization with Physics-Inspired Graph Neural Networks.
• DETR - Flax implementation of DETR: End-to-end Object Detection with Transformers using Sinkhorn solver and parallel bipartite matching.
• awesome-jax-flax-llms - Collection of LLMs implemented in JAX & Flax
• JAX RL - Implementations of reinforcement learning algorithms.

Haiku

• OGB Large-Scale Challenge - LSC).
• Persistent Evolution Strategies - Code used for the paper Unbiased Gradient Estimation in Unrolled Computation Graphs with Persistent Evolution Strategies.
• NuX - Normalizing flows with JAX.
• Two Player Auction Learning - JAX implementation of the paper Auction learning as a two-player game.
• Adversarial Robustness - Reference code for Uncovering the Limits of Adversarial Training against Norm-Bounded Adversarial Examples and Fixing Data Augmentation to Improve Adversarial Robustness.
• Gated Linear Networks - GLNs are a family of backpropagation-free neural networks.
• Glassy Dynamics - Open source implementation of the paper Unveiling the predictive power of static structure in glassy systems.
• Normalizer-Free Networks - Official Haiku implementation of NFNets.
• OGB-LSC - This repository contains DeepMind's entry to the PCQM4M-LSC (quantum chemistry) and MAG240M-LSC (academic graph)
• AlphaFold - Implementation of the inference pipeline of AlphaFold v2.0, presented in Highly accurate protein structure prediction with AlphaFold.

JAX

• Amortized Bayesian Optimization - Code related to Amortized Bayesian Optimization over Discrete Spaces.
• Accurate Quantized Training - Tools and libraries for running and analyzing neural network quantization experiments in JAX and Flax.
• BNN-HMC - Implementation for the paper What Are Bayesian Neural Network Posteriors Really Like?.
• JAX-DFT - 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 - Reference code for the paper A General and Adaptive Robust Loss Function.
• Symbolic Functionals - Demonstration from Evolving symbolic density functionals.
• TriMap - Official JAX implementation of TriMap: Large-scale Dimensionality Reduction Using Triplets.
• Fourier Feature Networks - Official implementation of Fourier Features Let Networks Learn High Frequency Functions in Low Dimensional Domains.
• kalman-jax - Approximate inference for Markov (i.e., temporal) Gaussian processes using iterated Kalman filtering and smoothing.
• jaxns - Nested sampling in JAX.

NumPyro

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

Trax

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

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

Tutorials and Blog Posts

Equinox

• Plugging Into JAX by Nick Doiron - Compares Flax, Haiku, and Objax on the Kaggle flower classification challenge.
• 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.
• 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.
• Deterministic ADVI in JAX by Martin Ingram - Walk through of implementing automatic differentiation variational inference (ADVI) easily and cleanly with JAX.
• Introduction to JAX by Kevin Murphy - Colab that introduces various aspects of the language and applies them to simple ML problems.
• 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.
• 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.
• Achieving 4000x Speedups with PureJaxRL - A blog post on how JAX can massively speedup RL training through vectorisation.
• Ensemble networks by Mat Kelcey - Ensemble nets are a method of representing an ensemble of models as one single logical model.
• Out of distribution (OOD) detection by Mat Kelcey - Implements different methods for OOD detection.
• Tutorial: image classification with JAX and Flax Linen by 8bitmp3 - Learn how to create a simple convolutional network with the Linen API by Flax and train it to recognize handwritten digits.
• Extending JAX with custom C++ and CUDA code by Dan Foreman-Mackey - Tutorial demonstrating the infrastructure required to provide custom ops in JAX.
• 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.
• 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.
• Get started with JAX by Aleksa Gordić - A series of notebooks and videos going from zero JAX knowledge to building neural networks in Haiku.
• 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.

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 via composable function transformations in Python | NeurIPS 2019 | Skye Wanderman-Milne - JAX intro presentation in Program Transformations for Machine Learning workshop.
• 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.
• JAX, Flax & Transformers 🤗 - 3 days of talks around JAX / Flax, Transformers, large-scale language modeling and other great topics.