JAX: grad, jit and vmap in Practice
JAX combines NumPy-style arrays with program transformations for machine learning and scientific computing. Start with pure functions, then transform them into gradients, compiled programs or batches.
- •Compose differentiation, compilation and batching.
- •Write a pure loss function and take a gradient step.
- •Separate compilation overhead from steady-state runtime.
Sources & further reading
Official JAX getting started ↗JAX installation guide ↗In-Depth Architecture & Mechanics
A transformation-first workflow
Write a function whose result depends only on its arguments. grad differentiates a scalar-valued function; jit compiles compatible computations; vmap maps a function across a batch dimension. Compose these transformations rather than maintaining separate loop implementations.
Tracing and shapes
JIT traces computations using abstract values. Python side effects and data-dependent Python branches do not behave like ordinary eager execution. Keep array shapes predictable, use JAX control-flow primitives when needed, and expect recompilation when shapes or static arguments change.
Practice
Run the example, print the initial and updated losses, then batch predictions with vmap. Warm up the compiled function before timing it and call block_until_ready on results to include asynchronous device execution. Compare correctness before comparing speed.
Randomness
Pass random keys explicitly and split them before independent random operations. Reusing a key repeats its random values. Check the official installation guide for your CPU, GPU or TPU environment.
import jax
import jax.numpy as jnp
x = jnp.array([1., 2., 3.])
y = 2 * x
def loss(w, x, y):
return jnp.mean((w * x - y) ** 2)
@jax.jit
def step(w, x, y):
return w - 0.05 * jax.grad(loss)(w, x, y)
w = jnp.array(0.)
updated = step(w, x, y)
updated.block_until_ready()
print(loss(w, x, y), loss(updated, x, y))
predict = jax.vmap(lambda value: updated * value)
print(predict(x))