JAX for PyTorch Users

Explain how JAX differs from NumPy and PyTorch in terms of execution model, immutability, and functional purity.

Create, manipulate, and inspect jax.numpy arrays, contrasting them with NumPy arrays and PyTorch tensors.

Use jax.device_put and inspect array device placement and dtypes on a single accelerator.

Rewrite a stateful Python/NumPy routine into a pure function suitable for JAX transformations.

Diagnose common pitfalls when porting NumPy code to JAX, such as in-place updates, Python control flow, and side effects.

Use .at[].set/add for functional in-place array updates.

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.

Define what a pytree is in JAX and identify pytree leaves vs container nodes.

Use jax.tree_util.tree_map, tree_leaves, tree_structure, and tree_flatten/unflatten in practice.

Apply transformations like vmap and grad to functions operating on pytrees of arrays.

Register a custom Python class as a pytree node using register_pytree_node_class.

Distinguish between static and dynamic fields when registering a custom pytree.

Use equinox.tree_at and filter utilities to surgically update fields of a pytree.

Diagnose common pytree errors such as treedef mismatches and leaf-type inconsistencies.

Distinguish forward-mode (jvp) and reverse-mode (vjp) autodiff and choose appropriately by problem shape.

Compute Jacobians and Hessians using jax.jacfwd, jax.jacrev, and jax.hessian.

Implement Hessian-vector products efficiently by composing jvp and grad.

Define a custom forward and backward rule using jax.custom_vjp.

Apply jax.lax.stop_gradient to block gradient flow at chosen points.

Explain the mechanism of gradient checkpointing (remat) for trading compute for memory in long backward passes.

Explain Equinox's philosophy: models as pytrees of arrays with no separate parameter dictionaries.

Define a custom eqx.Module subclass with parameters, submodules, and static fields.

Use eqx.filter, eqx.partition, and eqx.combine to separate trainable arrays from static data.

Apply eqx.filter_jit and eqx.filter_grad to train Equinox models cleanly.

Use an inference-mode flag to conditionally apply operations like dropout, demonstrating correct PRNG key management.

Manage and update evolving state (e.g., BatchNorm running statistics) by partitioning the model and returning updated state from the forward pass.