a function that computes the gradient of "f"), a
  vectorization/batching transform ("vmap(f)" returns a function that
  computes "f" over batches of inputs), and others.

These function transforms can compose with each other arbitrarily.
  For example, composing "vmap(grad(f))" computes a quantity called
  per-sample-gradients that stock PyTorch cannot efficiently compute
  today.

Why composable function transforms?
There are a number of use cases that are tricky to do in PyTorch
today:


computing per-sample-gradients (or other per-sample quantities)


running ensembles of models on a single machine


efficiently batching together tasks in the inner-loop of MAML


efficiently computing Jacobians and Hessians


efficiently computing batched Jacobians and Hessians


Composing "vmap()", "grad()", and "vjp()" transforms allows us to
express the above without designing a separate subsystem for each.
This idea of composable function transforms comes from the JAX