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?

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

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

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

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