import torch


def truncated_normal(tensor, mean=0, std=0.5):
    """ Truncated normal distribution, similar to tf.random.truncated_normal
    :param tensor: A tensor of arbitrary shape.
    :param mean: mean of the distribution
    :param std: std of the distribution

    :returns: A tensor with the same shape of x, with values draw from the designated truncated normal.
    """

    with torch.no_grad():
        size = tensor.shape
        tmp = tensor.new_empty(size + (8,)).normal_(mean=mean, std=std)
        valid = (tmp < 2 * std) & (tmp > -2 * std)
        ind = valid.max(-1, keepdim=True)[1]
        tmp = tmp.gather(-1, ind).squeeze(-1)
        return tmp