import torch


def ind2sub(shape, index):
    """
    A PyTorch implementation of MATLAB's "ind2sub" function

    Parameters
    ----------
    shape : torch.Size | list | tuple
        shape of the 2D matrix
    index : torch.Tensor
        (n,) linear indices

    Returns
    -------
    rows : torch.Tensor
        (n,) row subscripts
    cols : torch.Tensor
        (n,) column subscripts

    """
    # checks
    assert isinstance(shape, torch.Size) or \
           isinstance(shape, list) or \
           isinstance(shape, tuple)
    assert isinstance(index, torch.Tensor) and len(index.shape) == 1
    valid_index = index < shape[0]*shape[1]
    assert valid_index.all()
    if not len(shape) == 2:
        raise NotImplementedError('only implemented for 2D case.')
    # compute inds
    cols = index % shape[0]
    rows = index // shape[0]

    return rows, cols