NanoGPT (124M) quality in 7.8 8xH100-minutes

AlphaFold 3 inference pipeline.

MoVQGAN - model for the image encoding and reconstruction

Implementation of the specific Transformer architecture from PaLM - Scaling Language Modeling with Pathways - in Jax (Equinox framework)

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

The AI Scientist: Towards Fully Automated Open-Ended Scientific Discovery 🧑‍🔬

JAX implementation of the Mistral 7b v0.2 model

DeepMind's Tacotron-2 Tensorflow implementation

A multi-voice TTS system trained with an emphasis on quality

An Open-Ended Agentic Simulator

Jax implementation of VIT-VQGAN

Hybrid ML physics model of the Earth's atmosphere

Simple single-file baselines for Q-Learning in pure-GPU setting

Cost aware hyperparameter tuning algorithm

JAX-Toolbox

Simplifying reinforcement learning for complex game environments

🙌 OpenHands: Code Less, Make More

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

The release codes of LA-MCTS with its application to Neural Architecture Search.

Fast Diffusion Models with Transformers

Official codebase used to develop Vision Transformer, SigLIP, MLP-Mixer, LiT and more.

Train vision models using JAX and 🤗 transformers

tiny vision language model