import torch


def sample_concentrated_unit_ball(shape):
    """
    Samples from the unit ball non-uniformly, 
    with direction being sampled uniformly, 
    and radius being samples from Uniform[0,1]
    (non-uniformity is to avoid concentration on surface
    effects)
    
    Parameters
    ----------
    shape : List[int]
        shape of samples
    
    Returns
    -------
    torch.Tensor
        Samples

    """

    x = torch.randn(*shape)
    x /= torch.linalg.vector_norm(x,dim=-1,keepdim=True)
    u = torch.rand(*(shape[:-1] + (1,)))
    x *= u
    return x