import torch


def l_infinity_gradient(x, *, max_tol=0.) -> torch.Tensor:
    """
    Compute the gradient of the infinity-norm |x|∞
    Args:
      max_tol: If |x[i]| is within max_tol to |x|∞, then we consider the
      gradient w.r.t x[i].

    Return:
      gradient: A (num_possible_gradient x x_dim) size tensor. gradient[i] is
      the i'th possible gradient.
    """
    assert (isinstance(x, torch.Tensor))
    x_dim = x.shape[0]
    assert (x.shape == (x_dim, ))
    inf_norm = torch.norm(x, p=float("inf"))
    if inf_norm <= max_tol:
        return torch.cat((torch.eye(
            x_dim, dtype=x.dtype), -torch.eye(x_dim, dtype=x.dtype)),
                         dim=0)
    grad = torch.where(torch.abs(x - inf_norm) <= max_tol, 1, 0) + \
        torch.where(torch.abs(x + inf_norm) <= max_tol, -1, 0)
    nonzero_grad_indices = torch.nonzero(grad).squeeze(1).tolist()
    if len(nonzero_grad_indices) == 1:
        return grad.unsqueeze(0).type(x.dtype)
    all_grad = torch.zeros((len(nonzero_grad_indices), x_dim), dtype=x.dtype)
    for (i, grad_index) in enumerate(nonzero_grad_indices):
        all_grad[i, grad_index] = grad[grad_index]
    return all_grad