def index_select(self, dim: int, index: torch.Tensor):
        """Returns a sub-matrix selected according to the given index.

        Parameters
        ----------
        dim : int
            The dim to select from matrix, should be 0 or 1. `dim = 0` for
            rowwise selection and `dim = 1` for columnwise selection.
        index : torch.Tensor
            The selection index indicates which IDs from the `dim` should
            be chosen from the matrix.
            Note that duplicated ids are allowed.

        The function does not support autograd.

        Returns
        -------
        SparseMatrix
            The sub-matrix which contains selected rows or columns.

        Examples
        --------

        >>> indices = torch.tensor([0, 1, 1, 2, 3, 4], [0, 2, 4, 3, 5, 0]])
        >>> val = torch.tensor([0, 1, 2, 3, 4, 5])
        >>> A = dglsp.spmatrix(indices, val)

        Case 1: Select rows by IDs.

        >>> row_ids = torch.tensor([0, 1, 4])
        >>> A.index_select(0, row_ids)
        SparseMatrix(indices=tensor([[0, 1, 1, 2],
                                     [0, 2, 4, 0]]),
                     values=tensor([0, 1, 2, 5]),
                     shape=(3, 6), nnz=4)

        Case 2: Select columns by IDs.

        >>> column_ids = torch.tensor([0, 4, 5])
        >>> A.index_select(1, column_ids)
        SparseMatrix(indices=tensor([[0, 4, 1, 3],
                                     [0, 0, 1, 2]]),
                     values=tensor([0, 5, 2, 4]),
                     shape=(5, 3), nnz=4)
        """
        if dim not in (0, 1):
            raise ValueError("The selection dimension should be 0 or 1.")
        if isinstance(index, torch.Tensor):
            return SparseMatrix(self.c_sparse_matrix.index_select(dim, index))
        raise TypeError(f"{type(index).__name__} is unsupported input type.")
Returns a sub-matrix selected according to the given index.

        Parameters
        ----------
        dim : int
            The dim to select from matrix, should be 0 or 1. `dim = 0` for
            rowwise selection and `dim = 1` for columnwise selection.
        index : torch.Tensor
            The selection index indicates which IDs from the `dim` should
            be chosen from the matrix.
            Note that duplicated ids are allowed.

        The function does not support autograd.

        Returns
        -------
        SparseMatrix
            The sub-matrix which contains selected rows or columns.

        Examples
        --------

        >>> indices = torch.tensor([0, 1, 1, 2, 3, 4], [0, 2, 4, 3, 5, 0]])
        >>> val = torch.tensor([0, 1, 2, 3, 4, 5])
        >>> A = dglsp.spmatrix(indices, val)

        Case 1: Select rows by IDs.

        >>> row_ids = torch.tensor([0, 1, 4])
        >>> A.index_select(0, row_ids)
        SparseMatrix(indices=tensor([[0, 1, 1, 2],
                                     [0, 2, 4, 0]]),
                     values=tensor([0, 1, 2, 5]),
                     shape=(3, 6), nnz=4)

        Case 2: Select columns by IDs.

        >>> column_ids = torch.tensor([0, 4, 5])
        >>> A.index_select(1, column_ids)
        SparseMatrix(indices=tensor([[0, 4, 1, 3],
                                     [0, 0, 1, 2]]),
                     values=tensor([0, 5, 2, 4]),
                     shape=(5, 3), nnz=4)
        
python/dgl/sparse/sparse_matrix.py