JAX

HuggingFace Transformershuggingface/transformers

A model-definition framework for state-of-the-art machine learning models across text, vision, audio, and multimodal tasks.

funnlpfighting41love/funNLP

A comprehensive collection of Chinese NLP resources, datasets, tools, and pre-trained models for developers and researchers.

annotated_deep_learning_paper_implementationslabmlai/annotated_deep_learning_paper_implementations

A collection of 60+ annotated PyTorch implementations of deep learning papers with side-by-side explanatory notes.

sglangsgl-project/sglang

A high-performance serving framework for large language models and multimodal models, delivering low-latency and high-throughput inference.

AlphaFolddeepmind/alphafold

Open source implementation of AlphaFold 2, a deep learning system for highly accurate protein structure prediction.

Vision Transformergoogle-research/vision_transformer

Official JAX/Flax implementation of Vision Transformer (ViT) and MLP-Mixer for image recognition, with pre-trained models.

korniakornia/kornia

A differentiable computer vision library for PyTorch, providing geometric vision and image processing algorithms for AI workflows.

Dopaminegoogle/dopamine

A research framework for fast prototyping of reinforcement learning algorithms, designed for easy experimentation and reproducibility.

cleanrlvwxyzjn/cleanrl

A deep reinforcement learning library offering high-quality, single-file implementations of algorithms like PPO, DQN, and SAC for research and education.

AlphaFold3google-deepmind/alphafold3

AlphaFold 3 is an AI model that predicts the 3D structures of proteins and their interactions with other biomolecules like DNA, RNA, and ligands.

Traxgoogle/trax

An end-to-end deep learning library focused on clear code, speed, and research, built by Google Brain.

Reformergoogle/trax

An end-to-end deep learning library focused on clear code and speed, used for research and production by Google Brain.

Flaxgoogle/flax

A neural network library for JAX designed for flexibility, enabling researchers to experiment with new training forms by modifying training loops.

Flax NNXgoogle/flax

A neural network library for JAX designed for flexibility, enabling researchers to experiment with new training forms by modifying training loops.

DeepChemdeepchem/deepchem

An open-source Python library for applying deep learning to drug discovery, materials science, quantum chemistry, and biology.

Whisper JAXsanchit-gandhi/whisper-jax

A JAX implementation of OpenAI's Whisper model offering up to 70x faster transcription on TPUs.

Oryxtensorflow/probability

A library for probabilistic reasoning and statistical analysis integrated with TensorFlow and JAX.

Acmegoogle-deepmind/acme

A research framework for reinforcement learning providing modular building blocks and reference agent implementations.

Scenicgoogle-research/scenic

A JAX library for rapid prototyping of large-scale attention-based vision models across images, video, audio, and multimodal data.

DI-engineopendilab/DI-engine

A comprehensive Python-first reinforcement learning framework with modular abstractions for decision intelligence applications.

Haikudeepmind/dm-haiku

A JAX-based neural network library that provides a simple, object-oriented programming model for building and training models.

BRAXgoogle/brax

A fast, differentiable physics engine built with JAX for massively parallel rigid body simulation on accelerator hardware.

Equinoxpatrick-kidger/equinox

A JAX library for neural networks and scientific computing with PyTorch-like syntax and full ecosystem compatibility.

NumPyropyro-ppl/numpyro

A lightweight probabilistic programming library using NumPy and JAX for autograd and JIT compilation to GPU/TPU/CPU.

Mctxdeepmind/mctx

A JAX-native library implementing Monte Carlo tree search algorithms like AlphaZero and MuZero for reinforcement learning research.

EasyLMyoung-geng/EasyLM

A JAX/Flax-based framework for easy and scalable pre-training, fine-tuning, evaluation, and serving of large language models.

MaxTextgoogle/maxtext

A high-performance, scalable LLM library and reference implementation written in pure Python/JAX for training on TPUs and GPUs.

Neural Tangentsgoogle/neural-tangents

A high-level neural network API for specifying and analyzing infinite-width neural networks as Gaussian Processes in Python.

Objectrongoogle-research-datasets/Objectron

A large-scale dataset of object-centric video clips with 3D bounding box annotations and AR metadata for 3D object detection research.

Optaxdeepmind/optax

A gradient processing and optimization library for JAX, designed for research with composable building blocks.

Optaxgoogle-deepmind/optax

A gradient processing and optimization library for JAX, designed for research with composable building blocks.

cvxpylayerscvxgrp/cvxpylayers

A Python library for constructing differentiable convex optimization layers in PyTorch, JAX, and MLX using CVXPY.

Diffraxpatrick-kidger/diffrax

A JAX-based library providing numerical differential equation solvers for ODEs, SDEs, and CDEs with autodifferentiation and GPU support.

Penzaigoogle-deepmind/penzai

A JAX research toolkit for building, editing, and visualizing neural networks as legible, functional pytree data structures.

Big Transfer (BiT)google-research/big_transfer

Official repository for Big Transfer (BiT) models, providing pre-trained visual representations for efficient transfer learning across computer vision tasks.

Jraphgoogle-deepmind/jraph

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