awesome-jax

Haiku

Focused on simplicity, created by the authors of Sonnet at DeepMind

Objax

Has an object oriented design similar to PyTorch

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

Jraph

Lightweight graph neural network library

Neural Tangents

High-level API for specifying neural networks of both finite and width

HuggingFace

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

Levanter

Legible, Scalable, Reproducible Foundation Models with Named Tensors and JAX

EasyLM

LLMs made easy: Pre-training, finetuning, evaluating and serving LLMs in JAX/Flax

NumPyro

Probabilistic programming based on the Pyro library

Chex

Utilities to write and test reliable JAX code

Optax

Gradient processing and optimization library

RLax

Library for implementing reinforcement learning agents

JAX, M.D.

Accelerated, differential molecular dynamics

Coax

Turn RL papers into code, the easy way

Distrax

Reimplementation of TensorFlow Probability, containing probability distributions and bijectors

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

Equivariant MLP

Construct equivariant neural network layers

jax-resnet

Implementations and checkpoints for ResNet variants in Flax

Parallax

Immutable Torch Modules for JAX

jax-unirep

Library implementing the for protein machine learning applications

jax-flows

Normalizing flows in JAX

sklearn-jax-kernels

kernel matrices using JAX

jax-cosmo

Differentiable cosmology library

efax

Exponential Families in JAX

mpi4jax

Combine MPI operations with your Jax code on CPUs and GPUs

imax

Image augmentations and transformations

FlaxVision

Flax version of TorchVision

Oryx

Probabilistic programming language based on program transformations

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

PIX

PIX is an image processing library in JAX, for 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

Einshape

DSL-based reshaping library for JAX and other frameworks

ALX

Open-source library for distributed matrix factorization using Alternating Least Squares, more info in

Diffrax

Numerical differential equation solvers in JAX

tinygp

The of Gaussian process libraries in JAX

gymnax

Reinforcement Learning Environments with the well-known gym API

Mctx

Monte Carlo tree search algorithms in native JAX

KFAC-JAX

Second Order Optimization with Approximate Curvature for NNs

TF2JAX

Convert functions/graphs to JAX functions

jwave

A library for differentiable acoustic simulations

GPJax

Gaussian processes in JAX

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 )

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 🤗

Kernex

Differentiable stencil decorators in JAX

MaxText

A simple, performant and scalable Jax LLM written in pure Python/Jax and targeting Google Cloud TPUs

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

Vectorized board game environments for RL with an AlphaZero example

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

dynamiqs

High-performance and differentiable simulations of quantum systems with JAX

kalman-jax

Approximate inference for Markov (i.e., temporal) Gaussian processes using iterated Kalman filtering and smoothing

jaxns

Nested sampling in JAX

Amortized Bayesian Optimization

Code related to

Accurate Quantized Training

Tools and libraries for running and analyzing neural network quantization experiments in JAX and Flax

BNN-HMC

Implementation for the paper

JAX-DFT

One-dimensional density functional theory (DFT) in JAX, with implementation of

Robust Loss

Reference code for the paper

Symbolic Functionals

Demonstration from

TriMap

Official JAX implementation of

Performer

Flax implementation of the Performer (linear transformer via FAVOR+) architecture

JaxNeRF

Implementation of with multi-device GPU/TPU support

mip-NeRF

Official implementation of

RegNeRF

Official implementation of

Big Transfer (BiT)

Implementation of

JAX RL

Implementations of reinforcement learning algorithms

Implementation of

Minimal implementation of

Distributed Shampoo

Implementation of

NesT

Official implementation of

XMC-GAN

Official implementation of

FNet

Official implementation of

GFSA

Official implementation of

IPA-GNN

Official implementation of

Flax Models

Collection of models and methods implemented in Flax

Protein LM

Implements BERT and autoregressive models for proteins, as described in and

Slot Attention

Reference implementation for

Vision Transformer

Official implementation of

FID computation

Port of to Flax

ARDM

Official implementation of

D3PM

Official implementation of

Gumbel-max Causal Mechanisms

Code for , with extra code in

Latent Programmer

Code for the ICML 2021 paper

SNeRG

Official implementation of

Spin-weighted Spherical CNNs

Adaptation of

VDVAE

Adaptation of , original code at

MUSIQ

Checkpoints and model inference code for the ICCV 2021 paper

AQuaDem

Official implementation of

Combiner

Official implementation of

Dreamfields

Official implementation of the ICLR 2022 paper

GIFT

Official implementation of

Light Field Neural Rendering

Official implementation of

A JAX/Flax implementation of the Sharpened Cosine Similarity layer

GNNs for Solving Combinatorial Optimization Problems

A JAX + Flax implementation of

AlphaFold

Implementation of the inference pipeline of AlphaFold v2.0, presented in

Adversarial Robustness

Reference code for and

Bootstrap Your Own Latent

Implementation for the paper

Gated Linear Networks

GLNs are a family of backpropagation-free neural networks

Glassy Dynamics

Open source implementation of the paper

MMV

Code for the models in

Normalizer-Free Networks

Official Haiku implementation of

NuX

Normalizing flows with JAX

OGB-LSC

This repository contains DeepMind's entry to the (quantum chemistry) and (academic graph) tracks of the (OGB-LSC)

Persistent Evolution Strategies

Code used for the paper

Two Player Auction Learning

JAX implementation of the paper

WikiGraphs

Baseline code to reproduce results in

Reformer

Implementation of the Reformer (efficient transformer) architecture

lqg

Official implementation of Bayesian inverse optimal control for linear-quadratic Gaussian problems from the paper

Simple neural network from scratch in JAX

JAX's core design, how it's powering new research, and how you can start using it

Introduction to Bayesian modelling using NumPyro

JAX intro presentation in workshop

Presentation of TPU host access with demo

Tutorial created by Zico Kolter, David Duvenaud, and Matt Johnson with Colab notebooks avaliable in

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

JAX, Flax & Transformers 🤗

3 days of talks around JAX / Flax, Transformers, large-scale language modeling and other great topics

White paper describing an early version of JAX, detailing how computation is traced and compiled

Introduces JAX, M.D., a differentiable physics library which includes simulation environments, interaction potentials, neural networks, and more

Uses JAX's JIT and VMAP to achieve faster differentially private than existing libraries

White paper describing the XLB library: benchmarks, validations, and more details about the library

Neural network building blocks from scratch with the basic JAX operators

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

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

Compares Flax, Haiku, and Objax on the Kaggle flower classification challenge

Introduction to both JAX and Meta-Learning

Concise implementation of

Tutorial on implementing path tracing

Ensemble nets are a method of representing an ensemble of models as one single logical model

Implements different methods for OOD detection

Understand how autodiff works using JAX

Showcases how to go from a PyTorch-like style of coding to a more Functional-style of coding

Extending JAX with custom C++ and CUDA code by Dan Foreman-Mackey

Tutorial demonstrating the infrastructure required to provide custom ops in JAX

Explores how JAX can power the next generation of scalable neuroevolution algorithms

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

Walk through of implementing automatic differentiation variational inference (ADVI) easily and cleanly with JAX

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

Colab that introduces various aspects of the language and applies them to simple ML problems

Tutorial on the different ways to write an MCMC sampler in JAX along with speed benchmarks

Tutorial on how to add a progress bar to compiled loops in JAX using the module

Get started with JAX by Aleksa Gordić

A series of notebooks and videos going from zero JAX knowledge to building neural networks in Haiku

A tutorial on writing a simple end-to-end training and evaluation pipeline in JAX, Flax and Optax

A tutorial on 3D volumetric rendering of scenes represented by Neural Radiance Fields in JAX

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

A blog post on how JAX can massively speedup RL training through vectorisation