JAX Pathways

103 available lessons22 phases30 hands-on projects

Learn JAX.
Build working models.

From arrays and gradients to training, evaluation, and deployment. Learn the ideas, run the experiments, and build with an AI tutor alongside you.

Created by Taha Bouhsine. Run on your own machine.

Start with basic Python and a laptop CPU. The course is in active development.

Start the courseExplore learning pathsStar on GitHub0starsStart with an AI tutor

Learn with your own AI tutor

Start with a placement check, choose a learning path, and work through explanations and exercises with your agent.

1. Download the course workspace, unzip it, and open the jaxpathways folder in your agent. 2. Run this in that folder’s terminal (requires Node.js/npm):

npx skills add . --skill start-learning learn-jax jax-course-guide check-jax

3. Ask your agent: Use start-learning to help me learn JAX from this workspace.

Codex · Claude Code · Cursor · other skill-enabled agents

Using a browser chatbot, or need setup help?

For ChatGPT or another browser chatbot, download the first lesson, attach it to your chat, and ask for a guided walkthrough. The command above installs skills for coding agents that can read local files. Read the setup guide.

Transform a pure function with grad for derivatives, vmap for batching, and jit for compilation.
One function. Three ways to transform it.gradvmapjit

Where do you want to take it?

Start with the same foundations. Then follow the work that interests you.

Explore all 7 career routes ↗

Training

Build training loops, neural networks and small Transformers.

19 available · 0 plannedLessons in the focus phases. Shared preparation is listed in the path.

Learning path ↗Career skills ↗

Inference

Take a trained model toward export, serving and deployment.

8 available · 0 plannedLessons in the focus phases. Shared preparation is listed in the path.

Learning path ↗Career skills ↗

Research

Use differentiation for scientific models and numerical experiments.

4 available · 0 plannedLessons in the focus phases. Shared preparation is listed in the path.

Learning path ↗Career skills ↗

Systems

Explore profiling, distributed computation and accelerator operations.

31 available · 0 plannedLessons in the focus phases. Shared preparation is listed in the path.

Learning path ↗Career skills ↗

A look inside the course.

Read an explanation, make a prediction, then run and change the code. A few lessons from the journey:

View the full curriculum ↗
  1. 00.01Set up your learning workspacePrepare your workspace and run a first example.40 min
  2. 01.01From NumPy to jax.numpyTranslate familiar NumPy operations into JAX.45 min
  3. 02.01Your first gradientDifferentiate a function and check the result.55 min
  4. 02.03Batch a function with vmapApply one function across a batch of inputs.60 min
  5. 02.04Compile a function with jitCompile a function and understand its first call.55 min
  6. 05.03A compiled train and evaluation stepConnect training and evaluation in a model workflow.95 min
available lesson drafts
103
integration projects
30
curriculum phases
22

All 103 lesson exercises run on CPU; planned capstones and accelerator validation are labeled in each path.

Put the pieces together.

Start with these two projects, then explore the specialist projects as you progress.

Inputs pass through a hidden layer to a classifier output

Build a neural classifier

Connect a network, an optimizer, and an evaluation step. Investigate what happens when training goes wrong.

Open the classifier project ↗
Browse all projects

Staged CPU practice with exercises and checks. Each project states its prerequisites.

Take the course with you.

Read away from the browser. Keep the explanations beside your editor, or start with the first lesson online.

Download EPUBPrintable reader

Book contents, author and edition details. The printable reader can be saved as a PDF from your browser. Read the first lesson online ↗