import torch


def proj(x, epsilon=0.00001):
    """
    Projects a matrix of vectors in hyperbolic geometry back into the unit ball
    Parameters
    ----------
    x
        the matrix with each row representing a vector
    epsilon
        the small difference to subtract to ensure that vectors remain in the unit ball

    Returns
    -------
        a matrix of vectors modified to ensure that all vectors are contained
        in the unit ball

    """
    size = x.size()
    norms = torch.norm(x, p=2, dim=1)
    #print(norms)
    norms = norms.repeat(1, size[1])
    res = x / norms
    res -= epsilon
    return res