How to Learn JAX, from jax.numpy to Sharded Training

How do you learn JAX?

You learn JAX by learning its rules for arrays and randomness first and then its transformations one at a time, writing a small function for each and checking that it gives the right answer. TensorTonic's JAX for ML Systems study plan has 35 coding problems in nine sections, starting with how JAX arrays and random keys work and ending with the decoding loops and sharding patterns used to train and serve large models. Each problem is checked against tests, so you know a function is correct before you build on it. This guide gives a ten-step learning order, says what to practise at each step, compares JAX with PyTorch and lists the mistakes that slow most beginners down.

What is JAX, and how is it different from NumPy and PyTorch?

JAX is a Python library for numerical computing that turns ordinary array functions into compiled, differentiated, vectorized and sharded programs. Its jax.numpy module looks almost the same as NumPy, so a lot of NumPy code runs after changing the import. The differences come from one design choice, which is that JAX transforms functions, and a transformation only works reliably on a pure function whose output depends on its inputs alone.

That choice explains the rules that surprise people coming from NumPy or PyTorch. JAX arrays are immutable, so an update returns a new array instead of changing the old one. Random numbers come from an explicit key that you pass into each function, because a hidden global random state would be a side effect. A model is a function that takes its parameters as an argument, which is different from a PyTorch module that stores its parameters as attributes.

The four transformations you will use most are jax.grad, which returns a function that computes gradients, jax.jit, which compiles a function with the XLA compiler, jax.vmap, which adds a batch dimension, and the sharding tools that spread a computation across devices. They compose, so a jitted function can take the gradient of a vmapped loss, and that composability is the main reason people choose JAX.

What do you need before learning JAX?

You need working Python, comfortable NumPy and a clear picture of how gradient descent trains a model before learning JAX. Knowing how attention works in a Transformer helps once you reach the later sections of the study plan, which build inference and decoding code.

• For Python, you should be able to write functions, closures and nested data structures without looking things up, and this checklist of the Python you need for ML lets you test yourself.
• For NumPy, you should understand array shapes, broadcasting, indexing and reductions, which the NumPy study plan covers in 25 problems across five sections.
• For gradients, you should know what a backward pass computes, and the Build Micrograd from Scratch study plan has 25 problems that build a small autograd engine so that jax.grad stops feeling like magic.
• For attention, the lesson on the attention mechanism covers what the batched attention and decoding problems assume.

You do not need to know PyTorch first. If you already know it, the guide on how to learn PyTorch is a useful reference for where the two libraries line up, because tensors, losses, optimizers and training loops mean the same thing in both.

In what order should you learn JAX?

You should learn JAX in the order jax.numpy, grad, jit and tracing, vmap, PRNG keys, pytrees, a training loop with Optax-style updates, control flow with lax, sharding, and finally profiling and memory. The early steps teach the rules every transformation depends on, and the later steps combine them into the code that trains and serves real models. The table pairs each step with a lesson to read where one exists and the part of the JAX study plan that practises it.

The study plan groups its problems by topic, so a few steps above draw on more than one section. Random keys appear in the first section and again in sampling, and lax.scan appears in the jit section and again in decoding.

How do jax.numpy arrays differ from NumPy arrays?

jax.numpy arrays differ from NumPy arrays in three main ways, which are that they cannot be changed in place, that they default to 32-bit floats and integers, and that they live on a device such as a CPU, GPU or TPU. To change one element you write x.at[i].set(v), and the same .at syntax supports add, multiply, min and max, each returning a new array. JAX keeps 64-bit types switched off unless you enable the jax_enable_x64 option, so code that relies on float64 precision gives slightly different numbers.

How does jax.grad work?

jax.grad works by taking a function that returns a scalar and returning a new function that computes the gradient of that scalar with respect to the first argument. The argnums option picks other arguments, and jax.value_and_grad returns the loss and the gradient together so you do not run the forward pass twice. When the loss function also returns extra values such as accuracy, has_aux=True tells JAX to pass them through without differentiating them.

Because gradients are functions, you can transform them again. Combining jax.jvp with jax.grad gives a Hessian-vector product without building the full Hessian, and wrapping a per-example loss in vmap before taking the gradient gives one gradient per example, which is how per-example clipping works. jax.custom_vjp lets you replace the backward pass of a function, which is how a straight-through estimator passes gradients through a rounding step.

What happens when jax.jit traces a function?

When jax.jit traces a function, it calls the Python code once with abstract tracer values that carry only a shape and a dtype, records the array operations, and compiles that record into a single XLA program. Later calls with the same shapes, dtypes and static arguments reuse the compiled program and skip Python entirely. A call with a new shape triggers a fresh trace and a fresh compile.

Tracing explains most jit errors. A Python if statement on a traced value fails because the tracer has no concrete value to compare, and a Python print inside the function runs once during tracing instead of on every call. Arguments that change the structure of the program, such as a layer count or a mode flag, should be marked with static_argnums or static_argnames so that JAX treats them as compile-time constants.

What does jax.vmap do?

jax.vmap takes a function written for a single example and returns a function that runs it over a batch, adding the batch dimension to every operation inside. The in_axes argument says which axis of each input to map over, and None means an input is shared across the batch, which is how you evaluate one input against an ensemble of parameter sets. Nesting two vmaps maps over two axes, so a distance function for two points becomes a pairwise distance matrix.

Why does JAX use explicit random keys?

JAX uses explicit random keys because a hidden global random state would be a side effect, and side effects break jit, vmap and parallel execution. You create a key with jax.random.key(0), split it with jax.random.split into fresh subkeys, and pass one subkey to each random call such as jax.random.normal. The same key always gives the same numbers, which makes runs reproducible across devices and means you must never use one key twice for numbers that should be independent.

What is a pytree in JAX?

A pytree in JAX is any nested structure of lists, tuples, dicts and registered classes whose leaves are arrays, and JAX transformations accept and return pytrees as easily as single arrays. Model parameters are usually a dict of layer dicts, and jax.grad returns gradients in exactly the same structure. jax.tree.map applies a function to every leaf, so an optimizer update for the whole model is one line, and computing a global gradient norm means reducing over jax.tree.leaves.

What does a JAX training loop look like?

A JAX training loop calls one jitted step function for every batch, and that function takes the parameters, the optimizer state and the batch and returns new parameters, new optimizer state and the loss. Nothing is updated in place, so the loop simply rebinds its variables to the returned values. The Optax library uses this pattern, with an optimizer that has an init function and an update function, and writing the same two functions yourself is the best way to understand it.

1. Initialize the parameters as a pytree using a random key, and create the optimizer state with optimizer.init(params).
2. Write a loss function that takes the parameters and a batch and returns a scalar.
3. Inside the step function, call jax.value_and_grad on the loss to get the loss and the gradients.
4. Call optimizer.update(grads, opt_state, params) to get the updates and the new optimizer state, then optax.apply_updates(params, updates) to get the new parameters.
5. Wrap the step function in jax.jit and call it in a Python loop over batches, passing the returned parameters and state into the next call.

Gradient accumulation fits the same pattern. A lax.scan over micro-batches sums the gradient pytrees, and the step applies one update with the average, which the lesson on gradient accumulation in LLM training explains.

How does control flow work with jax.lax?

Control flow inside jitted code works through jax.lax primitives, which take functions as arguments so that JAX can trace every branch and every loop body. jnp.where picks values element by element, lax.cond and lax.switch pick one branch based on a traced value, lax.fori_loop and lax.while_loop run a loop without unrolling it, and lax.scan carries a state through a sequence and stacks the outputs. A Python for loop over a fixed number of steps still works under jit, but it unrolls into one copy of the body per step, which makes long loops slow to compile.

lax.scan is the one to learn well, because recurrent layers, decoding loops and gradient accumulation all use it. Its body takes a carry and one slice of the input and returns a new carry and one slice of the output, and the shapes of the carry must stay the same on every step.

How does sharding work in JAX?

Sharding in JAX works by arranging devices into a named mesh, describing how each array is split across that mesh with a PartitionSpec, and letting jit compile a program that runs on all the devices at once. jax.make_mesh builds the mesh, a NamedSharding pairs the mesh with a PartitionSpec such as one that splits the batch axis over a data axis, and jax.device_put places an array with that sharding. The compiler then inserts the communication the computation needs.

When you want to write the per-device code and the communication yourself, shard_map runs a function on each device's shard and lets you call collectives such as lax.psum, lax.all_gather and lax.ppermute over a named axis. pmap is the older API for the same idea and still appears in many tutorials, while new code usually uses jit with shardings or shard_map. Data parallelism averages gradients with a psum over the data axis, and tensor parallelism splits a weight matrix across devices so each one computes part of a layer.

How do you profile JAX code?

You profile JAX code by separating compile time from run time first, because the first call to a jitted function includes tracing and compiling and later calls do not. JAX also dispatches work asynchronously, so a timer stopped right after a call measures only how long it took to queue the work, and calling block_until_ready on the result waits for the computation to finish. For a full picture, jax.profiler.trace records a trace you can open in TensorBoard or Perfetto to see which operations take the time.

Memory is the other half of performance. jax.checkpoint, also called remat, recomputes activations in the backward pass instead of storing them, and bfloat16 halves the memory of activations and weights. The lesson on mixed precision training explains when lower precision is safe and why loss scaling is needed for float16.

Should you learn JAX or PyTorch?

Most people should learn PyTorch first and JAX second, because PyTorch is used more widely in jobs, courses and open-source model code. JAX is worth learning early if you plan to work with TPUs, on a team that already uses JAX, or on research that needs custom gradients, many vectorized runs or large sharded training jobs.

The ideas transfer in both directions. Tensors, broadcasting, losses, optimizers and training loops are the same in both libraries, so the second one is much quicker to learn than the first. JAX makes you think about functions, explicit state and compilation from the start, and that habit also helps when you later use torch.compile or write kernels. The guides on how to learn PyTorch and how to learn Triton cover the other two paths, and Triton is the right next step if you want to write the GPU kernels that compilers like XLA generate for you.

What should you build after learning the basics of JAX?

After learning the basics of JAX, you should build the pieces of a language model's inference path, because they use every transformation together and have clear right answers to test against. The later sections of the JAX for ML Systems study plan cover rotary embeddings, KV cache updates, top-k and top-p sampling, a grouped-query decode step, greedy decoding with scan, paged KV cache gathers, beam search and speculative decoding acceptance.

The lessons on rotary position embeddings, the KV cache, paged attention and speculative decoding explain each idea before you implement it. After that, the Inference Engineering study plan applies the same patterns to model serving, and the guide on inference engineering lays out the full build order. For the training side, Account for ZeRO Memory by Stage shows how sharding optimizer state, gradients and parameters changes the memory each device needs.

What mistakes do beginners make when learning JAX?

The mistakes beginners make most often in JAX are side effects inside jitted functions, Python control flow on traced values, reused random keys, recompilation caused by changing shapes, and attempts to update arrays in place. Each one comes from forgetting that jit traces a pure function once and then reuses the compiled result.

• A print call or an append to a global list inside a jitted function runs only while JAX traces it, so it appears once with tracer values and then never again. Use jax.debug.print to see values at run time, and return anything you want to keep as an output of the function.
• An if statement or a while condition on a traced value raises a TracerBoolConversionError, because the tracer has no concrete value during tracing. Replace it with jnp.where, lax.cond or lax.while_loop, or mark the argument as static when it really is a configuration value.
• Passing the same key to two random calls gives the same random numbers twice, which silently correlates dropout masks, initial weights or samples. Split the key every time you need new randomness and never reuse a key after splitting it.
• Calling a jitted function with a new input shape forces a new trace and compile, so variable-length batches or sequences can make every step slow. Pad inputs to a small set of fixed bucket lengths and use a mask, and turn on the jax_log_compiles option to see when compiles happen.
• Writing x[i] = v on a JAX array raises a TypeError because JAX arrays are immutable. Write x = x.at[i].set(v) instead, and remember to keep the returned array, since the original is unchanged. Inside jit, XLA can usually perform the update in place, so this costs nothing in compiled code.

Practice this

Then work through all 35 problems in the JAX for ML Systems study plan, starting with the Arrays the JAX way section, or compare the same ideas in PyTorch with the PyTorch study plan.

Related guides

• How to learn PyTorch, from tensors to a training loop
• How to learn Triton, from your first kernel to FlashAttention
• How to learn CUDA: a path from C to fused kernels
• Inference engineering: what it is and how to learn it
• How to practise ML with no GPU, no dataset and no team