Skip to main content
Ctrl+K
JAX AI Stack - Home JAX AI Stack - Home
    • Get Started
    • Tutorials
    • Blog
    • News
  • Community & Events
  • Videos
  • .md
    • Get Started
    • Tutorials
    • Blog
    • News
  • Community & Events
  • Videos
  • .md
JAX AI Stack - Home JAX AI Stack - Home
  • JAX AI Stack

Getting started

  • Installing the stack
  • Getting started with JAX for ML
    • Part 1: JAX neural net basics
    • Part 2: Debug a variational autoencoder (VAE)
    • Part 3: Train a diffusion model for image generation

Tutorials

  • Visualize JAX model metrics with TensorBoard
  • Introduction to Data Loaders
    • Introduction to Data Loaders on CPU with JAX
    • Introduction to Data Loaders on GPU with JAX
  • From PyTorch to JAX
    • JAX for PyTorch users
    • Porting a PyTorch model to JAX

Example applications

  • Train a miniGPT language model with JAX
  • Basic text classification with 1D CNN
  • Text classification with a transformer language model using JAX
  • Machine Translation with encoder-decoder transformer model
  • Image segmentation with UNETR model
  • Image Captioning with Vision Transformer (ViT) model
  • Train a Vision Transformer (ViT) for image classification with JAX
  • Time series classification with CNN

Developer resources

  • Contribute to documentation

Other resources

  • Blog
  • News
  • Events
  • Videos
  • Learning resources
  • .md

Blog

Blog#

  • How Lightricks trains video diffusion models at scale with JAX on TPU - November 11, 2025 by Yaki Bitterman & Yoav HaCohen, PhD

  • How JAX makes high-performance economics accessible - November 11, 2025 by John Stachurski & Srikanth Kilaru

  • AI Innovators: How JAX on TPU is helping Escalante advance AI-driven protein design - September 23, 2025 by Srikanth Kilaru

  • Beyond backpropagation: JAX’s symbolic power unlocks new frontiers in scientific computing - September 9, 2025 by Srikanth Kilaru, Zekun Shi, Min Lin

  • An efficient path to production AI: Kakao’s journey with JAX and Cloud TPUs - August 19, 2025 by Minho Ryu, Nayeon Kim, Srikanth Kilaru

  • A roboticist’s journey with JAX: Finding efficiency in optimal control and simulation - July 29, 2025 by Srikanth Kilaru, Max Muchen Sun

  • Stanford’s Marin foundation model: The first fully open model developed using JAX - July 16, 2025 by Srikanth Kilaru, David Hall

previous

Contribute to documentation

next

News

By JAX team

© Copyright 2024, JAX team.