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()".

