JAX Quest

Tracks

1. 01🧬Why JAX Exists and What Makes It Different

   0/6 lessons

   Why JAX exists, what makes it different from PyTorch and TensorFlow

   Where JAX came from, why it isn't just another deep-learning framework, and what its functional-transform thesis means for the way you write numerics.

   Lesson list (6)

   • 01
     The Origin Story~8 min · origins, jax, tutorial
   • 02
     The Core Thesis: Composable Function Transformations~10 min · origins, jax, tutorial
   • 03
     JAX vs PyTorch vs TensorFlow~12 min · origins, jax, tutorial
   • 04
     NumPy on Accelerators and XLA Compilation~10 min · origins, jax, tutorial
   • 05
     JAX in the Ecosystem and Installation~10 min · origins, jax, tutorial
   • 06
     Your First JAX Program: Putting It All Together~10 min · origins, jax, tutorial

   ✓Quiz · 3 questions→

2. 02🔢JAX as NumPy — The Foundation

   0/6 lessons

   JAX as accelerated, differentiable NumPy

   Drop-in NumPy semantics on accelerators. Where jax.numpy mirrors NumPy, where it deliberately doesn't, and the dtype defaults that bite every newcomer.

   Lesson list (6)

   • 01
     jax.numpy: The NumPy-Compatible API~8 min · numpy, jax, tutorial
   • 02
     Dtypes: float32, bfloat16, and Why They Matter~10 min · numpy, jax, tutorial
   • 03
     Indexing, Slicing, and Math Operations~9 min · numpy, jax, tutorial
   • 04
     Reshaping Arrays~8 min · numpy, jax, tutorial
   • 05
     Where JAX Arrays Differ from NumPy~10 min · numpy, jax, tutorial
   • 06
     Practical Example: Porting NumPy to JAX~11 min · numpy, jax, tutorial

   ✓Quiz · 2 questions→

3. 03✨Pure Functions — The Rule That Unlocks Everything

   0/5 lessons

   Pure functions — the rule that unlocks every transformation

   Why JAX makes you write pure functions, what 'pure' actually means in trace-time terms, and how to refactor real impure NumPy code into JAX-clean shape.

   Lesson list (5)

   • 01
     What Pure Means and Why JAX Demands It~9 min · purity, jax, tutorial
   • 02
     What Breaks Purity: The Rogues Gallery~10 min · purity, jax, tutorial
   • 03
     Functional Updates: From Modify to Return~8 min · purity, jax, tutorial
   • 04
     Debugging Impure Code: Error Messages Decoded~10 min · purity, jax, tutorial
   • 05
     Practical Example: Refactoring Impure Code~11 min · purity, jax, tutorial

   ✓Quiz · 2 questions→

4. 04⚡jit — Just-In-Time Compilation with XLA

   0/5 lessons

   Just-In-Time compilation with XLA

   How jit traces, compiles, and caches. Static arguments, abstract values, and the canonical errors you'll trip over before you internalize tracing.

   Lesson list (5)

   • 01
     What jit Does: Tracing, Compiling, Caching~8 min · jit, jax, tutorial
   • 02
     Tracing Deep Dive: Abstract Values and ShapedArray~8 min · jit, jax, tutorial
   • 03
     Static Arguments and Control Flow~9 min · jit, jax, tutorial
   • 04
     Common jit Errors and Fixes~8 min · jit, jax, tutorial
   • 05
     Benchmarking JIT: Measuring Real Speedups~9 min · jit, jax, tutorial

   ✓Quiz · 2 questions→

5. 05📐grad — Automatic Differentiation, Functional Style

   0/5 lessons

   Automatic differentiation, functional style

   Functional auto-diff in JAX: grad, value_and_grad, has_aux, higher-order gradients, Jacobians, and the canonical compose-grad-with-jit-and-vmap pattern.

   Lesson list (5)

   • 01
     What grad Does: Functions In, Gradient Functions Out~8 min · grad, jax, tutorial
   • 02
     argnums, value_and_grad, and has_aux~9 min · grad, jax, tutorial
   • 03
     Higher-Order Gradients and Jacobians~8 min · grad, jax, tutorial
   • 04
     Composing grad with jit and vmap~9 min · grad, jax, tutorial
   • 05
     Practical Example: Training Linear Regression from Scratch~11 min · grad, jax, tutorial

   ✓Quiz · 2 questions→

6. 06🔄vmap — Automatic Vectorization

   0/4 lessons

   Automatic vectorization without batch loops

   Vectorization without writing a batch loop. in_axes, out_axes, per-example gradients, nested vmap — the transformation that ends batch dimension hand-coding.

   Lesson list (4)

   • 01
     What vmap Does: One Example to Batch~8 min · vmap, jax, tutorial
   • 02
     in_axes and out_axes: Controlling Vectorization~8 min · vmap, jax, tutorial
   • 03
     vmap + grad: Per-Example Gradients~9 min · vmap, jax, tutorial
   • 04
     Nested vmap and Advanced Patterns~9 min · vmap, jax, tutorial

   ✓Quiz · 2 questions→

7. 07🌐pmap — Multi-Device Parallelism

   0/5 lessons

   Multi-device SPMD parallelism

   SPMD parallelism over multiple devices. Collectives, distributed training pattern, and why jax.sharding + Mesh is the modern replacement for pure pmap.

   Lesson list (5)

   • 01
     What pmap Does: SPMD Across Devices~8 min · pmap, jax, tutorial
   • 02
     Collective Operations: Communicating Between Devices~8 min · pmap, jax, tutorial
   • 03
     pmap + grad: Distributed Training Pattern~8 min · pmap, jax, tutorial
   • 04
     Modern Alternatives: jax.sharding and Mesh~9 min · pmap, jax, tutorial
   • 05
     Practical Example: Data-Parallel Training Pattern~12 min · pmap, jax, tutorial

   ✓Quiz · 2 questions→

8. 08🎲Random Numbers — The Functional Way

   0/6 lessons

   Functional PRNG keys and reproducible randomness

   Why JAX makes randomness explicit, how PRNG keys split and partition cleanly, and the patterns that stop you from accidentally training with the same seed forever.

   Lesson list (6)

   • 01
     Why JAX Handles Randomness Differently~8 min · random, jax, tutorial
   • 02
     PRNG Keys and Splitting~8 min · random, jax, tutorial
   • 03
     Key Management Patterns~9 min · random, jax, tutorial
   • 04
     Random Distributions in JAX~8 min · random, jax, tutorial
   • 05
     Randomness with vmap and the Partitionable PRNG~11 min · random, jax, tutorial
   • 06
     Common Mistakes and Reproducibility~13 min · random, jax, tutorial

   ✓Quiz · 2 questions→

9. 09🌳Pytrees — JAX's Universal Data Structure

   0/5 lessons

   JAX's universal nested-structure abstraction

   Pytrees are JAX's nested-structure abstraction — every parameter pack, every state dict, every model. Map, flatten, register custom nodes.

   Lesson list (5)

   • 01
     What Is a Pytree?~8 min · pytrees, jax, tutorial
   • 02
     jax.tree.map — Transforming Pytrees~8 min · pytrees, jax, tutorial
   • 03
     Pytree Utilities: flatten, unflatten, and Debugging~10 min · pytrees, jax, tutorial
   • 04
     Custom Pytree Nodes~11 min · pytrees, jax, tutorial
   • 05
     Pytrees in Practice: Neural Networks as Pytrees~12 min · pytrees, jax, tutorial

   ✓Quiz · 2 questions→

10. 10🧠Building Neural Networks — Flax, Equinox, and Friends

    0/6 lessons

    Flax NNX, Equinox, and JAX-native neural networks

    Why core JAX has no nn.Module, what fills that gap (Flax NNX, Equinox), and how the same CNN looks in both libraries side by side. Plus a Transformer block.

    Lesson list (6)

    • 01
      Why JAX Doesn't Include a Neural Network Library~8 min · neural-nets, jax, tutorial
    • 02
      Flax NNX — The New Standard~9 min · neural-nets, jax, tutorial
    • 03
      Equinox — Models as Pytrees~10 min · neural-nets, jax, tutorial
    • 04
      Common Layers and Building Blocks~9 min · neural-nets, jax, tutorial
    • 05
      Building a CNN: Side-by-Side Comparison~11 min · neural-nets, jax, tutorial
    • 06
      Building a Transformer Block~12 min · neural-nets, jax, tutorial

    ✓Quiz · 2 questions→

11. 11🔄Training Loops — Optax and the JAX Way

    0/5 lessons

    Optax, schedules, and explicit training loops

    Optax composes gradient transformations like Lego. Schedules, the canonical training step, and how to scan + checkpoint with Orbax for robust runs.

    Lesson list (5)

    • 01
      Why JAX Training Loops Are Explicit~8 min · training, jax, tutorial
    • 02
      Optax: Composable Gradient Transformations~8 min · training, jax, tutorial
    • 03
      Learning Rate Schedules~8 min · training, jax, tutorial
    • 04
      The Complete Training Loop~9 min · training, jax, tutorial
    • 05
      Scan for Efficient Training and Checkpointing with Orbax~11 min · training, jax, tutorial

    ✓Quiz · 2 questions→

12. 12⚡Advanced Transformations and Patterns

    0/5 lessons

    custom_vjp, remat, scan, and meta-learning patterns

    Custom gradient rules, gradient checkpointing for memory, scan for structured control flow under jit, jax.debug, and meta-learning grad-of-grad patterns.

    Lesson list (5)

    • 01
      Custom Gradient Rules: custom_vjp and custom_jvp~8 min · advanced, jax, tutorial
    • 02
      Gradient Checkpointing with jax.checkpoint / jax.remat~10 min · advanced, jax, tutorial
    • 03
      Scan Patterns and Structured Control Flow~9 min · advanced, jax, tutorial
    • 04
      Debugging Under JIT and Sharding~11 min · advanced, jax, tutorial
    • 05
      Meta-Learning: Grad-of-Grad Patterns~9 min · advanced, jax, tutorial

    ✓Quiz · 2 questions→

13. 13🔬JAX for Scientific Computing Beyond ML

    0/5 lessons

    JAX beyond ML — Diffrax, NumPyro, differentiable physics

    JAX outside ML: differential equations with Diffrax, Bayesian inference with NumPyro, physics simulators, and a fully differentiable simulation example.

    Lesson list (5)

    • 01
      Why JAX Excels at Scientific Computing~8 min · scientific, jax, tutorial
    • 02
      Differential Equations with Diffrax~9 min · scientific, jax, tutorial
    • 03
      Bayesian Inference with NumPyro~8 min · scientific, jax, tutorial
    • 04
      Physics Simulations and the JAX Ecosystem~10 min · scientific, jax, tutorial
    • 05
      Practical Example: Differentiable Physics Simulation~10 min · scientific, jax, tutorial

    ✓Quiz · 2 questions→

14. 14🚀Ecosystem, Deployment, and What's Next

    0/5 lessons

    Production deployment, profiling, and the road ahead

    Exporting models for production, the TPU/GPU runtime story, profiling tools, Pallas custom kernels, and an honest 2026 take on where to go next.

    Lesson list (5)

    • 01
      Exporting JAX Models for Production~8 min · ecosystem, jax, tutorial
    • 02
      JAX on TPUs and GPUs~8 min · ecosystem, jax, tutorial
    • 03
      Profiling and Debugging Performance~8 min · ecosystem, jax, tutorial
    • 04
      The JAX AI Stack and Pallas Custom Kernels~8 min · ecosystem, jax, tutorial
    • 05
      JAX vs PyTorch in 2026 and Where to Go Next~12 min · ecosystem, jax, tutorial

    ✓Quiz · 2 questions→