dimension of the inputs should be mapped over. "in_dims"
        should have a structure like the inputs. If the "in_dim" for a
        particular input is None, then that indicates there is no map
        dimension. Default: 0.
  * **out_dims** (*int** or **Tuple**[**int**]*) -- Specifies
    where the mapped dimension should appear in the outputs. If
    "out_dims" is a Tuple, then it should have one element per
    output. Default: 0.

  * **randomness** (*str*) -- Specifies whether the randomness in
    this vmap should be the same or different across batches. If
    'different', the randomness for each batch will be different.
    If 'same', the randomness will be the same across batches. If
    'error', any calls to random functions will error. Default:
    'error'. WARNING: this flag only applies to random PyTorch
    operations and does not apply to Python's random module or
    numpy randomness.
