A playbook for effectively prompting post-trained LLMs

Matrix-free linear algebra in JAX.

Entropy Based Sampling and Parallel CoT Decoding

Frechet inception distance (FID) evaluation in JAX

An implementation of the Llama architecture, to instruct and delight

Access large language models from the command-line

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

🧱 Modula software package

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