Table of Contents
Shortcuts

# torch.autograd.Function.vmap

static Function.vmap(info, in_dims, *args)

  Defines a rule for the behavior of this autograd.Function underneath `torch.vmap()`. For a `torch.autograd.Function()` to support `torch.vmap()`, you must either override this staticmethod, or set `generate_vmap_rule` to `True` (you may not do both).

  If you choose to override this staticmethod: it must accept

  - an `info` object as the first argument. `info.batch_size` specifies the size of the dimension being vmapped over, while `info.randomness` is the randomness option passed to `torch.vmap()`.

  - an `in_dims` tuple as the second argument. For each arg in `args`, `in_dims` has a corresponding `Optional[int]`. It is `None` if the arg is not a Tensor or if the arg is not being vmapped over, otherwise, it is an integer specifying what dimension of the Tensor is being vmapped over.

  - `*args`, which is the same as the args to `forward()`.

  The return of the vmap staticmethod is a tuple of `(output, out_dims)`. Similar to `in_dims`, `out_dims` should be of the same structure as `output` and contain one `out_dim` per output that specifies if the output has the vmapped dimension and what index it is in.

  Please see Extending torch.func with autograd.Function for more details.

## 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.