Skip to content

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.