def sqrt_gaussian()

in noise.py [0:0]

• 14 lines of code
• 5 McCabe index (conditional complexity)

def sqrt_gaussian(mean, precision, device=None):
    """
    Samples from a square-root Gaussian distribution with the specified
    DxN_1xN_2x...xN_n  tensor `mean` and tensor `precision` which should be a
    scalar or have shape N_1xN_2x...xN_n. The zero-mean square-root Gaussian
    distribution is a multi-variate distribution probability density function:

    p(x) \propto exp(-precision * ||x||_2) with x \in \mathbb{R}^D

    Each column in the output corresponds to a single D-dimensional sample.
    """  # noqa: W905
    assert isinstance(precision, float) or precision.nelement() == 1 \
        or precision.shape == mean.shape[1:], "Invalid shape for precision."

    # sample directions as unit-norm Gaussian vectors:
    direction = torch.randn(mean.size(), device=mean.device)  # samples from Gaussian
    norm = torch.norm(direction, dim=0, keepdim=True)         # norm of the samples
    direction.div_(norm)                                      # unit-norm samples

    # sample norms from Gamma distribution:
    shape = torch.ones(mean.shape[1:]).mul_(mean.size(0))
    norm = torch.distributions.gamma.Gamma(shape, precision).sample()
    if norm.device != direction.device:
        norm = norm.to(direction.device)

    # return final sample:
    sample = direction.mul_(norm).add_(mean)
    if device is not None:
        sample = sample.to(device=device)
    return sample