JAX Tutorial: From Basics to Advanced JAX 教程：从入门到进阶

About this tutorial

A hands-on JAX tutorial in two parts, from the basics to advanced topics. It follows two courses on apxml.com — Getting Started with JAX and Advanced JAX: Performance, Optimization, and Scale — and adds five extra chapters on topics you run into in real projects: Pallas kernels, debugging, export and deployment, random numbers and sharding, and scientific computing.

Every chapter is a Jupyter notebook that was run end to end on JAX 0.12 (CPU), with all outputs included, so you can read it here without running anything. Key claims are backed by asserts in the code, and outdated APIs are shown side by side with the current way to write them.

How to use it

• Read online: open any chapter below.
• Run it yourself: on each chapter page, click Open in Colab. The first cell installs the tested JAX version; use the default CPU runtime. You can also download the .ipynb and run it locally.
• Exercises: each exercise has a ✍️ answer cell followed by a 🧪 self-check. ✅ means your answer passes; ⏳ means you haven’t defined the required name yet. Reference solutions are in each chapter’s appendix.
• First time through: follow the core learning path — read the main text and the 🟢 Must-read extras, and skip the ⚪ Optional ones for now. Looking for a specific API? Try the keyword index.

Part 1 · Basics

Part 2 · Advanced

Course explanations are reorganized from the key points of the original apxml.com courses, with a link to the original page at the start of each section. The original course content is © ApX Machine Learning.