Table of Contents
Shortcuts

# torch.func API Reference

## 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`.
`linearize` | Returns the value of `func` at `primals` and linear approximation at `primals`.
`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.
`functionalize` | functionalize is a transform that can be used to remove (intermediate) mutations and aliasing from a function, while preserving the function's semantics.

## 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 `functional_call()` is for: it accepts an nn.Module, the transformed `parameters`, and the inputs to the Module’s forward pass. It returns the value of running the Module’s forward pass with the replaced parameters.

Here’s how we would compute the Jacobian over the parameters

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

    def f(params, x):
        return torch.func.functional_call(model, params, x)

    x = torch.randn(3)
    jacobian = jacrev(f)(dict(model.named_parameters()), x)

`functional_call` | Performs a functional call on the module by replacing the module parameters and buffers with the provided ones.
`stack_module_state` | Prepares a list of torch.nn.Modules for ensembling with `vmap()`.
`replace_all_batch_norm_modules_` | In place updates `root` by setting the `running_mean` and `running_var` to be None and setting track_running_stats to be False for any nn.BatchNorm module in `root`

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

- Patching Batch Norm

## Docs

Access comprehensive developer documentation for PyTorch

View Docs

## Tutorials

Get in-depth tutorials for beginners and advanced developers

View Tutorials

## Resources

Find development resources and get your questions answered

View Resources

To analyze traffic and optimize your experience, we serve cookies on this site. By clicking or navigating, you agree to allow our usage of cookies. As the current maintainers of this site, Facebook’s Cookies Policy applies. Learn more, including about available controls: Cookies Policy.