import torch


def least_squares(a, b, rcond=None):
    """
    A PyTorch implementation of NumPy's "linalg.lstsq" function

    Parameters
    ----------
    a : torch.Tensor
        (m,n) "Coefficient" matrix
    b : torch.Tensor
        (m,) or (m,k) "dependent variable" values
    rcond : float
        Cutt-off ratio for small singular values of a. For the
        purposes of rank determination, singular values are treated
        as zero if they are smaller than rcond times the largest
        singular value of a. If `None`, the default will use the
        machine precision times max(m, n)

    Returns
    -------
    x : torch.Tensor
        (n,) or (n,k) least-squares solution
    residuals: torch.Tensor
        (1,) or (k,) or (0,) sums of residuals; squared Euclidean 2-norm for
        each column in b - a*x. If rank(a) < n or m <= n, this is an empty
        array. If b is 1-dimensional, this is a (1,) shape array. Otherwise
        the shape is (k,).
    rank : int
        rank of matrix a
    s : torch.Tensor
        (min(m,n),) singular values of matrix a

    """
    m,n = a.shape
    if rcond is None:
        rcond = max(a.shape)*torch.finfo(a.dtype).eps
    U, s, V = torch.svd(a)
    rank = torch.sum(s > rcond*s[0]).item()
    s_inv = torch.where(s > rcond*s[0], s.reciprocal(), torch.zeros_like(s))
    x = torch.matmul(V, s_inv.unsqueeze(1) * torch.matmul(U.T, b))
    if rank < n or m <= n:
        residuals = torch.tensor([])
    else:
        residuals = torch.sum((a@x - b)**2, 0, keepdim=len(b.shape)==1)

    return x, residuals, rank, s