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.