Inference code for Llama models

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more

A high-throughput and memory-efficient inference and serving engine for LLMs

DSPy: The framework for programming—not prompting—foundation models

Open source code for AlphaFold.

Flax is a neural network library for JAX that is designed for flexibility.

SGLang is yet another fast serving framework for large language models and vision language models.

Google DeepMind's software stack for physics-based simulation and Reinforcement Learning environments, using MuJoCo.

JAX-based neural network library

Large language models (LLMs) made easy, EasyLM is a one stop solution for pre-training, finetuning, evaluating and serving LLMs in JAX/Flax.

Monte Carlo tree search in JAX

Probabilistic programming with NumPy powered by JAX for autograd and JIT compilation to GPU/TPU/CPU.

An Extensible Deep Learning Library

A JAX research toolkit for building, editing, and visualizing neural networks.

Optax is a gradient processing and optimization library for JAX.

Concurrent data pipelines in Python >>>

A simple, performant and scalable Jax LLM!

Numerical differential equation solvers in JAX. Autodifferentiable and GPU-capable. https://docs.kidger.site/diffrax/

A Graph Neural Network Library in Jax

A Python-level JIT compiler designed to make unmodified PyTorch programs faster.

Hardware accelerated, batchable and differentiable optimizers in JAX.

BlackJAX is a Bayesian Inference library designed for ease of use, speed and modularity.

An implementation of the Fermionic Neural Network for ab-initio electronic structure calculations

RL Environments in JAX 🌍

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

PIX is an image processing library in JAX, for JAX.