Introduction to JAX
PyTorch mutates tensors. TensorFlow builds graphs. JAX compiles pure functions. That last one changes how you think about deep learning. Write pure-function neural network code using JAX's functional API (jax.numpy, jax.grad, jax.jit, jax.vmap). Explain the key design difference between PyTorch's eager mutation and JAX's functional compilation model. Apply jit compilation and vmap vectorization to accelerate training loops compared to naive Python. Train a simple network in JAX and contrast the explicit state management with PyTorch's object-oriented approach. You know how to build neural networks in PyTorch. You define an nn.Module, call .backward(), step the optimizer. It works. Millions of people use it. But PyTorch has a constraint baked into its DNA: it traces operations eagerly, one at a time, in Python. Every tensor + tensor is a separate kernel launch. Every training step re-interprets the same Python code. This works fine until you need to train a 540-billion-parameter model across 2,048 TPUs. Then the overhead kills you. Google DeepMind trains Gemini on JAX. Anthropic trained Claude on JAX. These are not small operations -- they are the largest neural network training runs on Earth. They chose JAX because it treats your training loop as a compilable program, not a sequence of Python calls. JAX is NumPy with three superpowers: automatic differentiation, JIT compilation to XLA, and automatic vectorization. You write…
Introduction to JAX: PyTorch mutates tensors. TensorFlow builds graphs. JAX compiles pure functions. That last one changes how you think about deep learning.
This free lesson is part of the AI Engineering from Scratch curriculum. Read the full explanation, run the lesson code, and verify the result in the interactive reader or from the repository source.
Browse the complete course catalog or open this lesson on GitHub.