torch.func.vmap
torch.func.vmap(func, in_dims=0, out_dims=0, randomness='error', *, chunk_size=None)
vmap is the vectorizing map; "vmap(func)" returns a new function
   that maps "func" over some dimension of the inputs. Semantically,
   vmap pushes the map into PyTorch operations called by "func",
   effectively vectorizing those operations.
vmap is useful for handling batch dimensions: one can write a
   function "func" that runs on examples and then lift it to a
   function that can take batches of examples with "vmap(func)". vmap
   can also be used to compute batched gradients when composed with
   autograd.
Note:
 "torch.vmap()" is aliased to "torch.func.vmap()" for convenience.
 Use whichever one you'd like.

Parameters:
      * func (function) -- A Python function that takes one or
        more arguments. Must return one or more Tensors.
  * **in_dims** (*int** or **nested structure*) -- Specifies which
