Bonsai: Jax Implementations of Popular Models

Bonsai is a minimal, lightweight JAX implementation of popular models.

We're committed to making popular models accessible in JAX through simple, hackable, and concise code. Our aim is to lower the barrier to entry for JAX and promote academic innovation.

• LLM (Large Language Models): Qwen 3,
• dLLM (diffusion-based Large Language Models): (Coming soon) Llada, ...
• ASR (Automatic Speech Recognition): (Coming soon) Whisper, ...
• Image segmentation: SAM2, ...
• Computational Biology: (Coming soon) ESM, ...
• WFM (World Foundation Model): (Coming soon) Cosmos, ...

Got models you'd like to see in JAX? Add a request or contribute.

To get started with JAX Bonsai, follow these steps to set up your development environment and run the models.

Clone the JAX Bonsai repository to your local machine.

Install the latest repository.

Jump right into our Qwen3 model, implemented in 300 lines of code in JAX.

We welcome contributions! If you're interested in adding new models, improving existing implementations, or enhancing documentation, please see our Contributing Guidelines.

Join our discord to socialize with other JAX enthusiasts.

• JAX: Learn more about JAX, a super fast NumPy-based ML framework with automatic differentiation.
• The JAX ecosystem: Unlock unparalleled speed and scale for your next-generation models. Explore an incredible suite of tools and libraries that effortlessly extend JAX's capabilities, transforming how you build, train, and deploy.
• MaxText and MaxDiffusion: Industury solution for highly scalable, high-performant JAX model library via Google Cloud Platform.
• JAX LLM Examples: Example high-performant implementation of LLMs in pure JAX.