Transforms: grad and vmap
Transforms are compiler features, not library macros.
Use grad for reverse-mode differentiation over supported programs.
grad(loss_fn)(grad {} (var {} loss_fn))Use vmap to vectorize a per-example function over a batch dimension.
vmap(per_example_fn)(vmap {} (var {} per_example_fn) (d-name {} batch))The exact supported surface is still defined by the compiler and current tests; prefer small working programs and compiler feedback over extrapolating from planned future surface.