static Function.vmap(info, in_dims, *args)[source]

Define the behavior for this autograd.Function underneath torch.vmap().

For a torch.autograd.Function() to support torch.vmap(), you must either override this static method, 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.


Access comprehensive developer documentation for PyTorch

View Docs


Get in-depth tutorials for beginners and advanced developers

View Tutorials


Find development resources and get your questions answered

View Resources