import torch


def _solve_lstsq_subcols(rhs, lhs, lhs_col_idx):
    """Solves an over-determined linear system for selected LHS columns.
        A batched version of `torch.lstsq`.
    Args:
        rhs: right-hand side vectors
        lhs: left-hand side matrices
        lhs_col_idx: a slice of columns in lhs
    Returns:
        a least-squares solution for lhs * X = rhs
    """
    lhs = lhs.index_select(-1, torch.tensor(lhs_col_idx, device=lhs.device).long())
    return torch.matmul(torch.pinverse(lhs), rhs[:, :, None])