import torch


def get_distributions(attention: torch.FloatTensor):
    """Generates the distribution tensor

    Args:
        attention: The attention tensor.

    Returns:
        torch.Tensor: The distribution tensor.

    """
    dist = torch.ones_like(attention) * attention
    dist = dist / dist.sum(dim=1).view(-1, 1).repeat(1, attention.shape[1])
    dist.requires_grad = False
    return dist