import torch


def project_onto_l1_ball(x, eps=1):
    """
    Compute Euclidean projection onto the L1 ball for a batch.
    
      min ||x - u||_2 s.t. ||u||_1 <= eps
    
    Inspired by the corresponding numpy version by Adrien Gaidon.
    
    Parameters
    ----------
    x: (batch_size, *) torch array
      batch of arbitrary-size tensors to project, possibly on GPU
      
    eps: float
      radius of l-1 ball to project onto
    
    Returns
    -------
    u: (batch_size, *) torch array
      batch of projected tensors, reshaped to match the original
    
    Notes
    -----
    The complexity of this algorithm is in O(dlogd) as it involves sorting x.
    
    References
    ----------
    [1] Efficient Projections onto the l1-Ball for Learning in High Dimensions
        John Duchi, Shai Shalev-Shwartz, Yoram Singer, and Tushar Chandra.
        International Conference on Machine Learning (ICML 2008)
    """
    original_shape = x.shape
    x = x.view(x.shape[0], -1)
    x_norm = x - torch.mean(x, dim=0)
    x_norm = x_norm / torch.norm(x_norm)
    mask = (torch.norm(x_norm, p=1, dim=1) < eps).float().unsqueeze(1)
    mu, _ = torch.sort(torch.abs(x_norm), dim=1, descending=True)
    cumsum = torch.cumsum(mu, dim=1).float()
    arange = torch.arange(1, x_norm.shape[1] + 1, device=x.device)
    rho, _ = torch.max((mu * arange.float() > (cumsum - eps)).float() * arange.float(), dim=1)
    theta = (cumsum[torch.arange(x.shape[0]), (rho.cpu() - 1).long()] - eps) / rho
    proj = (torch.abs(x) - theta.unsqueeze(1)).clamp(min=0)
    x = mask * x_norm + (1 - mask) * proj * torch.sign(x_norm)
    return x.view(original_shape)