5 / 8 JAX & XLA Mental Model

Stacking the Transforms

You want the gradient of a per-example scalar loss with respect to inputs, applied across an entire batch, compiled into one efficient executable. Which composition is idiomatic JAX?

A

jax.jit(loss_fn, static_argnums=(0,))

B

jax.jit(jax.vmap(jax.grad(loss_fn)))

C

jax.grad(jax.jit(jax.vmap(loss_fn)))

D

jax.vmap(jax.jit(loss_fn), in_axes=None)