JAX

Jraphgoogle-deepmind/jraph

A lightweight library for building and training graph neural networks using JAX.

RLaxdeepmind/rlax

A JAX-based library providing reinforcement learning building blocks for implementing agents, supporting both on-policy and off-policy learning.

JAX, M.D.google/jax-md

A differentiable, hardware-accelerated molecular dynamics simulation framework built on JAX for computational physics and materials science.

BlackJAXblackjax-devs/blackjax

A fast, modular Bayesian inference library for JAX, providing composable samplers for CPU and GPU.

purejaxrlluchris429/purejaxrl

High-performance, end-to-end reinforcement learning implementations fully written in JAX for massive parallelization on GPUs.

SKRLToni-SM/skrl

A modular Python library for Reinforcement Learning with support for PyTorch, JAX, NVIDIA Warp, and multiple environment interfaces.

KerasCVkeras-team/keras-cv

A library of modular computer vision components built on Keras 3, supporting TensorFlow, JAX, and PyTorch backends.

JAXoptgoogle/jaxopt

Hardware-accelerated, batchable, and differentiable optimization algorithms implemented in JAX for machine learning research.

KerasNLPkeras-team/keras-nlp

A pretrained modeling library for Keras 3 offering simple, flexible, and fast access to models for text, image, and audio tasks.

Dynamaxprobml/dynamax

A Python library for probabilistic state space modeling and inference, built on JAX.

EvoJAXgoogle/evojax

A scalable, hardware-accelerated neuroevolution toolkit built on JAX for parallel training across TPUs/GPUs.

Chexdeepmind/chex

A library of utilities for writing and testing reliable JAX code, including assertions, debugging tools, and test variants.

mip-NeRFgoogle/mipnerf

Official JAX implementation of Mip-NeRF, a multiscale neural radiance field model for anti-aliased novel view synthesis.

gymnaxRobertTLange/gymnax

A JAX-based library providing accelerated reinforcement learning environments with full compatibility to the classic gym API.

Nucleotide Transformerinstadeepai/nucleotide-transformer

A collection of transformer-based foundation models for genomics and transcriptomics, enabling tasks like sequence analysis, functional prediction, and conversational DNA exploration.

Jumanjiinstadeepai/jumanji

A diverse suite of scalable reinforcement learning environments written in JAX for hardware-accelerated research.

Objaxgoogle/objax

An object-oriented machine learning framework built on JAX, designed for simplicity and readability in research.

evosaxRobertTLange/evosax

A comprehensive, high-performance library implementing 30+ Evolution Strategies in JAX for scalable optimization on modern hardware.

JAX RLikostrikov/jax-rl

JAX (Flax) implementations of reinforcement learning algorithms for continuous action spaces, designed for research.

OTT-JAXott-jax/ott

A JAX-powered library for solving large-scale optimal transport problems, including matching, barycenters, and neural approximations.

Levanterstanford-crfm/levanter

A JAX-based framework for training large language models with a focus on legibility, scalability, and reproducibility.

BrainPybrainpy/BrainPy

A flexible, efficient, and extensible JIT-compiled framework for computational neuroscience and brain-inspired computation.

NetKetnetket/netket

A Python library built on JAX for studying many-body quantum systems using neural networks and machine learning.

Distraxdeepmind/distrax

A JAX-native library of probability distributions and bijectors, reimplementing a subset of TensorFlow Probability with emphasis on readability and extensibility.

GPJaxthomaspinder/GPJax

A low-level Gaussian process framework in JAX and Flax, designed for maximum flexibility and close alignment with mathematical notation.

scikit-femkinnala/scikit-fem

A pure Python library for finite element assembly, transforming bilinear forms into sparse matrices and linear forms into vectors.

Pgxsotetsuk/pgx

A collection of GPU-accelerated parallel game simulators for reinforcement learning, built with JAX.

Optimistixpatrick-kidger/optimistix

A JAX library for nonlinear optimization including root finding, minimization, fixed points, and least squares.

Paxgoogle/paxml

A JAX-based machine learning framework for configuring and training large-scale models with high efficiency on TPUs and GPUs.

XLBAutodesk/XLB

A differentiable, massively parallel Lattice Boltzmann library in Python for physics-based machine learning and fluid dynamics simulations.

PIXdeepmind/dm_pix

An image processing library built on JAX, designed to be optimized and parallelized with JAX transformations.

Tequilaaspuru-guzik-group/tequila

A high-level Python framework for formulating, optimizing, and executing variational quantum algorithms on simulators and real hardware.

Extending JAX with custom C++ and CUDA code by Dan Foreman-Mackeydfm/extending-jax

A tutorial demonstrating how to extend JAX with custom C++ and CUDA operations for high-performance computing.

EasyDeLerfanzar/EasyDeL

A JAX-based framework for streamlined training, fine-tuning, and high-performance serving of large language and multimodal models.

QDaxadaptive-intelligent-robotics/QDax

A hardware-accelerated Python library for running Quality-Diversity and neuroevolution algorithms in minutes instead of days.

tinygpdfm/tinygp

An extremely lightweight Gaussian Process library for Python built on JAX with GPU acceleration and automatic differentiation.