def _power_law_distribution(size: int, exponent: float, device: torch.device) -> Tensor:
    """Returns a power law distribution summing up to 1.

    Args:
        size:
            The size of the distribution.
        exponent:
            The exponent for the power law distribution.
        device:
            The device to create tensor on.

    Returns:
        A power law distribution tensor summing up to 1.
    """
    k = torch.arange(1, size + 1, device=device)
    power_dist = torch.tensor(k ** (-exponent))
    power_dist = power_dist / power_dist.sum()
    return power_dist
Returns a power law distribution summing up to 1.

    Args:
        size:
            The size of the distribution.
        exponent:
            The exponent for the power law distribution.
        device:
            The device to create tensor on.

    Returns:
        A power law distribution tensor summing up to 1.
    
lightly/loss/pmsn_loss.py