Core Function Transformations: jit, grad, vmap

Apply jax.jit to compile a function and measure speedups versus eager execution.

Explain JAX's trace-time vs run-time distinction and what a Tracer object represents.

Predict and diagnose function retracing using static_argnums and jax.make_jaxpr.

Compute gradients of scalar-valued functions using jax.grad, including with has_aux and argnums.

Vectorize a per-example function across a batch using jax.vmap with in_axes/out_axes.

Compose jit, grad, and vmap correctly and reason about their order of application.

Use jax.lax.cond and jax.lax.switch to express conditional logic inside jit-compiled functions.

Use jax.lax.scan to implement efficient, parallelizable loops and recurrent patterns.