Libraries
24 projects
FlaxA neural network library for JAX designed for flexibility, enabling researchers to experiment with new training forms by modifying training loops.
Flax NNXA neural network library for JAX designed for flexibility, enabling researchers to experiment with new training forms by modifying training loops.
HaikuA JAX-based neural network library that provides a simple, object-oriented programming model for building and training models.
ObjaxAn object-oriented machine learning framework built on JAX, designed for simplicity and readability in research.
TraxAn end-to-end deep learning library focused on clear code, speed, and research, built by Google Brain.
JraphA lightweight library for building and training graph neural networks in JAX, providing graph data structures, utilities, and model implementations.
Neural TangentsA high-level neural network API for specifying and analyzing infinite-width neural networks as Gaussian Processes in Python.
HuggingFace TransformersA model-definition framework for state-of-the-art machine learning models across text, vision, audio, and multimodal tasks.
EquinoxA JAX library for neural networks and scientific computing with PyTorch-like syntax and full ecosystem compatibility.
ScenicA JAX library for rapid prototyping of large-scale attention-based vision models across images, video, audio, and multimodal data.
PenzaiA JAX research toolkit for building, editing, and visualizing neural networks as legible, functional pytree data structures.
LevanterA JAX-based framework for training large language models with a focus on legibility, scalability, and reproducibility.
EasyLMA JAX/Flax-based framework for easy and scalable pre-training, fine-tuning, evaluation, and serving of large language models.
NumPyroA lightweight probabilistic programming library using NumPy and JAX for autograd and JIT compilation to GPU/TPU/CPU.
ChexA library of utilities for writing and testing reliable JAX code, including assertions, debugging tools, and test variants.
OptaxA gradient processing and optimization library for JAX, designed for research with composable building blocks.
RLaxA JAX-based library providing reinforcement learning building blocks for implementing agents, supporting both on-policy and off-policy learning.
JAX, M.D.A differentiable, hardware-accelerated molecular dynamics simulation framework built on JAX for computational physics and materials science.
DistraxA JAX-native library of probability distributions and bijectors, reimplementing a subset of TensorFlow Probability with emphasis on readability and extensibility.
cvxpylayersA Python library for constructing differentiable convex optimization layers in PyTorch, JAX, and MLX using CVXPY.
NetKetA Python library built on JAX for studying many-body quantum systems using neural networks and machine learning.
BlackJAXA fast, modular Bayesian inference library for JAX, providing composable samplers for CPU and GPU.
DynamaxA Python library for probabilistic state space modeling and inference, built on JAX.
New Libraries
60 projectsFedJAXA JAX-based library for federated learning simulations that emphasizes ease-of-use in research.
Equivariant MLPA JAX library for automatically generating equivariant neural network layers for arbitrary symmetry groups via constraint solving.
jax-resnetFlax implementations and pretrained checkpoints for ResNet, Wide ResNet, ResNeXt, ResNet-D, and ResNeSt in JAX.
ParallaxA pure, immutable module system for JAX that replaces PyTorch-style imperative coding with declarative parameter trees.
OptimistixA JAX library for nonlinear optimization including root finding, minimization, fixed points, and least squares.
JAXoptHardware-accelerated, batchable, and differentiable optimization algorithms implemented in JAX for machine learning research.
jax-unirepA performant JAX reimplementation of the UniRep model for generating protein sequence representations.
flowjaxA JAX library for distributions, bijections, and normalizing flows implemented as Equinox modules.
jax-cosmoA differentiable cosmology library built with JAX for automatic differentiation of cosmological calculations.
OryxA library for probabilistic reasoning and statistical analysis integrated with TensorFlow and JAX.
Optimal Transport ToolsA Python toolbox for solving optimal transport problems with JAX-powered computational efficiency.
delta PVA photovoltaic simulator with automatic differentiation for solar cell modeling and optimization, built on JAX.
jaxlieA JAX library implementing Lie groups for rigid body transformations in computer vision and robotics.
BRAXA fast, differentiable physics engine built with JAX for massively parallel rigid body simulation on accelerator hardware.
flaxmodelsA collection of pretrained deep learning models (StyleGAN2, GPT2, VGG, ResNet) for the Jax/Flax ecosystem.
exojaxA JAX-based differentiable spectral modeling library for exoplanets, brown dwarfs, and M dwarfs.
PIXAn image processing library built on JAX, designed to be optimized and parallelized with JAX transformations.
bayexA lightweight Bayesian optimization library built on JAX for efficient optimization of expensive-to-evaluate functions.
JaxDFA JAX-based framework for building differentiable numerical simulators with arbitrary discretizations for physical systems.
tree-mathA Python library that enables mathematical operations on JAX pytrees, allowing numerical algorithms to work directly with complex nested data structures.
jax-modelsUnofficial JAX/Flax implementations of deep learning research papers for vision transformers and other architectures.
PGMaxA JAX-based library for loopy belief propagation on discrete factor graphs, enabling efficient probabilistic inference.
EvoJAXA scalable, hardware-accelerated neuroevolution toolkit built on JAX for parallel training across TPUs/GPUs.
evosaxA comprehensive, high-performance library implementing 30+ Evolution Strategies in JAX for scalable optimization on modern hardware.
SymJAXA symbolic programming library built on JAX for concise, explicit, and optimized machine learning computations.
mcxA JAX-powered probabilistic programming library focused on performant sampling methods for Bayesian inference on CPU, GPU, and TPU.
EinshapeA DSL-based library for unified tensor reshaping, squeezing, expanding, and transposing in JAX, TensorFlow, and NumPy.
ALXA collection of research code and datasets released by Google Research under open licenses.
DiffraxA JAX-based library providing numerical differential equation solvers for ODEs, SDEs, and CDEs with autodifferentiation and GPU support.
tinygpAn extremely lightweight Gaussian Process library for Python built on JAX with GPU acceleration and automatic differentiation.
gymnaxA JAX-based library providing accelerated reinforcement learning environments with full compatibility to the classic gym API.
MctxA JAX-native library implementing Monte Carlo tree search algorithms like AlphaZero and MuZero for reinforcement learning research.
KFAC-JAXA JAX library for second-order optimization of neural networks using the K-FAC curvature approximation algorithm.
TF2JAXAn experimental library that converts TensorFlow functions and graphs into JAX functions for reuse and fine-tuning within JAX codebases.
jwaveA JAX-based research framework for differentiable and parallelizable acoustic simulations, running on CPU, GPU, and TPU.
GPJaxA low-level Gaussian process framework in JAX and Flax, designed for maximum flexibility and close alignment with mathematical notation.
JumanjiA diverse suite of scalable reinforcement learning environments written in JAX for hardware-accelerated research.
EqxvisionA Python package providing popular computer vision model architectures built with Equinox for JAX.
SPUA provable, measurable secure computation device that enables privacy-preserving tensor operations using multi-party computation (MPC).
jax-tqdmAdd tqdm progress bars to JAX scans and loops using decorators, enabling side-effect-free progress tracking.
KernexKernex extends JAX with kmap and kscan for differentiable stencil computations, enabling efficient array transformations.
MaxTextA high-performance, scalable LLM library and reference implementation written in pure Python/JAX for training on TPUs and GPUs.
PaxA JAX-based machine learning framework for configuring and training large-scale models with high efficiency on TPUs and GPUs.
PraxisA layer library for JAX-based machine learning projects, optimized for large-scale ML.
purejaxrlHigh-performance, end-to-end reinforcement learning implementations fully written in JAX for massive parallelization on GPUs.
LoraxA JAX transform that implements LoRA (Low-Rank Adaptation) for efficient fine-tuning of large models with minimal memory overhead.
SCICOA Python package built on JAX for solving inverse problems in scientific imaging using optimization and prior models.
SpyxA compact spiking neural network library built on JAX and Haiku, offering high-performance training via surrogate gradient descent and neuroevolution.
BrainPyA flexible, efficient, and extensible JIT-compiled framework for computational neuroscience and brain-inspired computation.
OTT-JAXA JAX-powered library for solving large-scale optimal transport problems, including matching, barycenters, and neural approximations.
QDaxA hardware-accelerated Python library for running Quality-Diversity and neuroevolution algorithms in minutes instead of days.
JAX ToolboxA collection of CI pipelines, Docker images, and optimized examples to simplify JAX development on NVIDIA GPUs.
PgxA collection of GPU-accelerated parallel game simulators for reinforcement learning, built with JAX.
EasyDeLA JAX-based framework for streamlined training, fine-tuning, and high-performance serving of large language and multimodal models.
XLBA differentiable, massively parallel Lattice Boltzmann library in Python for physics-based machine learning and fluid dynamics simulations.
dynamiqsA Python library for GPU-accelerated and differentiable quantum systems simulation built with JAX.
NAVIXA JAX-powered reimplementation of MiniGrid offering over 1000x speedup for reinforcement learning experiments.
FDTDXAn efficient open-source Python package for 3D photonic nanostructure simulation and design using GPU-accelerated FDTD with automatic differentiation.
kvaxA FlashAttention 2 implementation for JAX with block-wise document mask optimization and context parallelism for efficient long-sequence training.
torchaxA PyTorch frontend for JAX that enables running PyTorch code on TPUs and provides seamless PyTorch-JAX interoperability.
JAX
4 projects
Fourier Feature NetworksA technique using Fourier feature mappings to enable neural networks to learn high-frequency functions in low-dimensional domains.
kalman-jaxA JAX-based framework for approximate inference in Markov Gaussian processes using iterated Kalman smoothing.
jaxnsA JAX-based probabilistic programming framework using nested sampling for fast Bayesian inference and evidence computation.
Amortized Bayesian OptimizationA collection of research code and datasets released by Google Research under open licenses.
Flax
8 projectsmip-NeRFOfficial JAX implementation of Mip-NeRF, a multiscale neural radiance field model for anti-aliased novel view synthesis.
Big Transfer (BiT)Official repository for Big Transfer (BiT) models, providing pre-trained visual representations for efficient transfer learning across computer vision tasks.
JAX RLJAX (Flax) implementations of reinforcement learning algorithms for continuous action spaces, designed for research.
NesTA vision transformer architecture that aggregates nested local transformers on image blocks for better accuracy, data efficiency, and convergence.
XMC-GANOfficial JAX implementation of XMC-GAN for text-to-image generation using cross-modal contrastive learning.
Vision TransformerOfficial JAX/Flax implementation of Vision Transformer (ViT) and MLP-Mixer for image recognition, with pre-trained models.
GNNs for Solving Combinatorial Optimization ProblemsA JAX + Flax implementation of physics-inspired graph neural networks for solving combinatorial optimization problems like Max-Cut and Maximum Independent Set.
Haiku
4 projectsAlphaFoldOpen source implementation of AlphaFold 2, a deep learning system for highly accurate protein structure prediction.
Adversarial RobustnessA collection of implementations and illustrative code accompanying DeepMind's published research papers across AI and machine learning.
Bootstrap Your Own LatentA collection of implementations and illustrative code accompanying DeepMind's published research papers across AI and scientific domains.
NuXA JAX library for building, training, and evaluating normalizing flows for probabilistic modeling.