import torch


def sample_generalized_normal(mean: torch.Tensor, lbd: float, p: int, shape):
    """Sample from a generalized Normal distribution.
    Modified according to:
    https://www.tensorflow.org/probability/api_docs/python/tfp/distributions/GeneralizedNormal

    Args:
        mean: Mean of the distribution.
        lbd: Parameter controlling the standard deviation of the distribution.
        p: Exponent of the distribution.
        shape: Shape of the samples to generate.
    """

    assert isinstance(lbd, float)

    ipower = 1.0 / p
    gamma_dist = torch.distributions.Gamma(ipower, 1.0)
    gamma_sample = gamma_dist.rsample(shape)
    # could speed up operations, but doesnt....
    # gamma_sample = torch._standard_gamma(torch.ones(shape) * ipower)
    binary_sample = torch.randint(low=0, high=2, size=shape, dtype=mean.dtype) * 2 - 1
    sampled = binary_sample * torch.pow(torch.abs(gamma_sample), ipower)
    return mean + lbd * sampled.to(mean.device)