Einsum-like high-level array sharding API for JAX

RL Environments in JAX 🌍

JetStream is a throughput and memory optimized engine for LLM inference on XLA devices, starting with TPUs (and GPUs in future -- PRs welcome).

TORAX: Tokamak transport simulation in JAX

Accelerated minigrid environments with JAX

A declarative drawing API in Python

Simple Transformer in Jax

Scalable neural net training via automatic normalization in the modular norm.

A massively parallel, optimal functional runtime in Rust

A massively parallel, high-level programming language

Universal Tensor Operations in Einstein-Inspired Notation for Python.

seqax = sequence modeling + JAX

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

Turn jitted jax functions back into python source code

Gaussian Markov Random Fields (GMRFs) and Integrated Nested Laplace Approximation (INLA)

A Jax-based library for designing and training transformer models from scratch.

(Crafter + NetHack) in JAX. ICML 2024 Spotlight.

Tensor Parallelism with JAX + Shard Map

Supercharge huggingface transformers with model parallelism.

S + Autograd + XLA :: S-parameter based frequency domain circuit simulations and optimizations using JAX.

A machine learning compiler for GPUs, CPUs, and ML accelerators

Inference code for Llama models

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

Development repository for the Triton language and compiler

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

State of the art inference for your bayesian models.

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