functorch¶

Function Transforms¶

`vmap`	vmap is the vectorizing map; `vmap(func)` returns a new function that maps `func` over some dimension of the inputs.
`grad`	`grad` operator helps computing gradients of `func` with respect to the input(s) specified by `argnums`.
`grad_and_value`	Returns a function to compute a tuple of the gradient and primal, or forward, computation.
`vjp`	Standing for the vector-Jacobian product, returns a tuple containing the results of `func` applied to `primals` and a function that, when given `cotangents`, computes the reverse-mode Jacobian of `func` with respect to `primals` times `cotangents`.
`jvp`	Standing for the Jacobian-vector product, returns a tuple containing the output of func(*primals) and the “Jacobian of `func` evaluated at `primals`” times `tangents`.
`jacrev`	Computes the Jacobian of `func` with respect to the arg(s) at index `argnum` using reverse mode autodiff
`jacfwd`	Computes the Jacobian of `func` with respect to the arg(s) at index `argnum` using forward-mode autodiff
`hessian`	Computes the Hessian of `func` with respect to the arg(s) at index `argnum` via a forward-over-reverse strategy.

Utilities for working with torch.nn.Modules¶

In general, you can transform over a function that calls a torch.nn.Module. For example, the following is an example of computing a jacobian of a function that takes three values and returns three values:

model = torch.nn.Linear(3, 3)

def f(x):
    return model(x)

x = torch.randn(3)
jacobian = jacrev(f)(x)
assert jacobian.shape == (3, 3)

However, if you want to do something like compute a jacobian over the parameters of the model, then there needs to be a way to construct a function where the parameters are the inputs to the function. That’s what make_functional() and make_functional_with_buffers() are for: given a torch.nn.Module, these return a new function that accepts parameters and the inputs to the Module’s forward pass.

`make_functional`	Given a `torch.nn.Module`, `make_functional()` extracts the state (params) and returns a functional version of the model, `func`.
`make_functional_with_buffers`	Given a `torch.nn.Module`, make_functional_with_buffers extracts the state (params and buffers) and returns a functional version of the model `func` that can be invoked like a function.
`combine_state_for_ensemble`	Prepares a list of torch.nn.Modules for ensembling with `vmap()`.

If you’re looking for information on fixing Batch Norm modules, please follow the guidance here

Patching Batch Norm