functorch API Reference¶
Function Transforms¶


Returns a function to compute a tuple of the gradient and primal, or forward, computation. 

Computes the Jacobian of 

vmap is the vectorizing map; 

Standing for the vectorJacobian product, returns a tuple containing the results of 
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.
Given a 

Given a 

Prepares a list of torch.nn.Modules for ensembling with 