numpy randomness.
  * **chunk_size** (*None** or **int*) -- If None (default), apply
    a single vmap over inputs. If not None, then compute the vmap
    "chunk_size" samples at a time. Note that "chunk_size=1" is
    equivalent to computing the vmap with a for-loop. If you run
    into memory issues computing the vmap, please try a non-None
    chunk_size.

Returns:
      Returns a new "batched" function. It takes the same inputs as
      "func", except each input has an extra dimension at the index
      specified by "in_dims". It takes returns the same outputs as
      "func", except each output has an extra dimension at the index
      specified by "out_dims".
Return type:
      Callable
One example of using "vmap()" is to compute batched dot products.
   PyTorch doesn't provide a batched "torch.dot" API; instead of
   unsuccessfully rummaging through docs, use "vmap()" to construct a
   new function.