Libraries
13 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.
TraxAn end-to-end deep learning library focused on clear code, speed, and research, built by Google Brain.
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.
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.
OptaxA gradient processing and optimization library for JAX, designed for research with composable building blocks.
cvxpylayersA Python library for constructing differentiable convex optimization layers in PyTorch, JAX, and MLX using CVXPY.
New Libraries
6 projects
OryxA library for probabilistic reasoning and statistical analysis integrated with TensorFlow and JAX.
BRAXA fast, differentiable physics engine built with JAX for massively parallel rigid body simulation on accelerator hardware.
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.
MctxA JAX-native library implementing Monte Carlo tree search algorithms like AlphaZero and MuZero for reinforcement learning research.
MaxTextA high-performance, scalable LLM library and reference implementation written in pure Python/JAX for training on TPUs and GPUs.
Flax
2 projectsBig Transfer (BiT)Official repository for Big Transfer (BiT) models, providing pre-trained visual representations for efficient transfer learning across computer vision tasks.
Vision TransformerOfficial JAX/Flax implementation of Vision Transformer (ViT) and MLP-Mixer for image recognition, with pre-trained models.
Haiku
3 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.