import torch


def sub2ind(shape, rows, cols):
    """
    A PyTorch implementation of MATLAB's "sub2ind" function

    Parameters
    ----------
    shape : torch.Size | list | tuple
        shape of the 2D matrix
    rows : torch.Tensor
        (n,) row subscripts
    cols : torch.Tensor
        (n,) column subscripts

    Returns
    -------
    index : torch.Tensor
        (n,) linear indices
    """
    # checks
    assert isinstance(shape, tuple) or isinstance(shape, list)
    assert isinstance(rows, torch.Tensor) and len(rows.shape) == 1
    assert isinstance(cols, torch.Tensor) and len(cols.shape) == 1
    assert len(rows) == len(cols)
    assert torch.all(rows < shape[0]) and torch.all(cols < shape[1])
    if not len(shape) == 2:
        raise NotImplementedError('only implemented for 2D case.')
    # compute inds
    ind_mat = torch.arange(shape[0]*shape[1]).view(shape)
    index = ind_mat[rows.long(), cols.long()]

    return index