import torch


def sample_diag_gauss(mean, std, generator=None, is_radial=False):
    """Get a sample from a multivariate Gaussian distribution with diagonal
    covariance matrix.

    Samples are produced using the reparametrization trick.

    Note: 
        If ``is_radial`` is set to ``True``, samples will instead by obtained
        from a radial BNN distribution instead of a multivariate Gaussian.
        For details refer to docstring of :func:`decode_and_sample_diag_gauss`.

    Args:
        mean: A list of tensors. See return value of method
            :func:`extract_mean_std`.
        std: A list of tensors with the same shapes as `mean`. See return value
            of method :func:`extract_mean_std`.
        generator (torch.Generator, optional): A generator can be passed to
            obtain more control over the reproducibility of the random sampling
            process.
        is_radial (bool, optional): If ``True``, the weights will be sampled
            according to a radial distribution, and not a Gaussian one.

    Returns:
        A list of tensors, where each is a sample from the diagonal Gaussian
        distributions defined by entries of `mean` and `std`.
    """
    sample = []
    device = mean[0].device if len(mean) > 0 else None
    for i, m in enumerate(mean):
        eps = torch.normal(torch.zeros_like(m), 1., generator=generator)

        if not is_radial:
            sample.append(m + eps * std[i])
        else:
            r = torch.normal(torch.tensor(0.).to(device), 1., \
                             generator=generator)
            sample.append(m +  std[i] * r * eps / torch.norm(eps, p=2))

    return sample