def _adapted_rsampling(
    shape: Union[Tuple[int, ...], torch.Size],
    dist: torch.distributions.Distribution,
    same_on_batch: Optional[bool] = False,
) -> Tensor:
    r"""Sample from a uniform reparameterized sampling function that accepts 'same_on_batch'.

    If same_on_batch is True, all values generated will be exactly same given a batch_size (shape[0]). By default,
    same_on_batch is set to False.
    """
    if isinstance(shape, tuple):
        shape = torch.Size(shape)

    if same_on_batch:
        rsample_size = torch.Size((1, *shape[1:]))
        rsample = dist.rsample(rsample_size)
        return rsample.repeat(shape[0], *[1] * (len(rsample.shape) - 1))
    return dist.rsample(shape)
Sample from a uniform reparameterized sampling function that accepts 'same_on_batch'.

    If same_on_batch is True, all values generated will be exactly same given a batch_size (shape[0]). By default,
    same_on_batch is set to False.
    
kornia/augmentation/utils/helpers.py