containing a struct matching the shape of the input:



f = lambda dict: torch.dot(dict['x'], dict['y'])
x, y = torch.randn(2, 5), torch.randn(5)
input = {'x': x, 'y': y}
batched_dot = torch.vmap(f, in_dims=({'x': 0, 'y': None},))
batched_dot(input)



By default, the output is batched along the first dimension.
   However, it can be batched along any dimension by using "out_dims"



f = lambda x: x ** 2
x = torch.randn(2, 5)
batched_pow = torch.vmap(f, out_dims=1)
batched_pow(x) # [5, 2]



For any function that uses kwargs, the returned function will not
   batch the kwargs but will accept kwargs



x = torch.randn([2, 5])
def fn(x, scale=4.):
  return x * scale
batched_pow = torch.vmap(fn)
assert torch.allclose(batched_pow(x), x * 4)
batched_pow(x, scale=x) # scale is not batched, output has shape [2, 2, 5]



Note:
 vmap does not provide general autobatching or handle variable-
