Fix description and not emplement error.

python/dgl/sparse/sparse_matrix.py

@@ -479,104 +479,125 @@ def is_diag(self):
         """Returns whether the sparse matrix is a diagonal matrix."""
         return self.c_sparse_matrix.is_diag()
 
-    def index_select(self, dim: int, index: torch.Tensor):
-        """Returns a sub-matrix selected according to the given index.
+    def rowwise_select(self, ids: torch.Tensor):
+        """Returns a row-wise sub-matrix selected by a list of row ids.
 
         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 : tensor.Tensor
-            The selection index indicates which IDs from the `dim` should
-            be chosen from the matrix.
-            Note that duplicated ids are allowed.
+        ids : tensor.Tensor
+            The row ids to select from matrix.
+            NOTE: The ids can include duplicate values.
 
-        The function does not support autograd.
+        The function has not supported autograd.
 
         Returns
         -------
         SparseMatrix
-            The sub-matrix which contains selected rows or columns.
+            The row-wise sub-matrix only contains the rows selected.
 
         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)
+        >>> A.rowwise_select(row_ids)
         SparseMatrix(indices=tensor([[0, 1, 1, 2],
                                      [0, 2, 4, 0]]),
                      values=tensor([0, 1, 2, 5]),
                      shape=(3, 6), nnz=4)
+        """
+        return self.c_sparse_matrix.rowwise_select(ids)
 
-        Case 2: Select columns by IDs.
+    def columnwise_select(self, ids: torch.Tensor):
+        """Returns a column-wise sub-matrix selected by ids.
 
+        Parameters
+        ----------
+        ids : tensor.Tensor
+            The column ids to select from matrix.
+            NOTE: The ids can include duplicate values.
+
+        The function has not supported autograd.
+
+        Returns
+        -------
+        SparseMatrix
+            The column-wise sub-matrix only contains the columns selected.
+
+        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)
         >>> 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]),
+        >>> A.columnwise_select(column_ids)
+        SparseMatrix(indices=tensor([[0, 1, 3, 4],
+                                     [0, 1, 2, 0]]),
+                     values=tensor([0, 2, 4, 6]),
                      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):
-            raise NotImplementedError
-        raise TypeError(f"{type(index).__name__} is unsupported input type.")
+        return self.c_sparse_matrix.columnwise_select(ids)
 
-    def range_select(self, dim: int, index: slice):
-        """Returns a sub-matrix selected according to the given range index.
+    def rowwise_select(self, start: int, end: int):
+        """Returns a row-wise sub-matrix selected by range [start, end).
 
         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 : slice
-            The selection slice indicates ID index from the `dim` should
-            be chosen from the matrix.
+        start : int
+            The start row to select from matrix.
+        end : int
+            The end row to select from matrix.
 
-        The function does not support autograd.
+        The function has not supported autograd.
 
         Returns
         -------
         SparseMatrix
-            The sub-matrix which contains selected rows or columns.
+            The row-wise sub-matrix between [start, end).
 
         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 with given slice object.
-
-        >>> A.range_select(0, slice(1, 3))
+        >>> A.rowwise_select(1, 3)
         SparseMatrix(indices=tensor([[0, 0, 1],
                                      [2, 4, 3]]),
                      values=tensor([1, 2, 3]),
                      shape=(2, 6), nnz=3)
+        """
+        return self.c_sparse_matrix.rowwise_select(start, end)
 
-        Case 2: Select columns with given slice object.
+    def columnwise_select(self, start: int, end: int):
+        """Returns a column-wise sub-matrix selected by range [start, end).
 
-        >>> A.range_select(1, slice(3, 6))
-        SparseMatrix(indices=tensor([[2, 1, 3],
-                                     [0, 1, 2]]),
-                     values=tensor([3, 2, 4]),
+        Parameters
+        ----------
+        start : int
+            The start column to select from matrix.
+        end : int
+            The end column to select from matrix.
+
+        The function has not supported autograd.
+
+        Returns
+        -------
+        SparseMatrix
+            The column-wise sub-matrix between [start, end).
+
+        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)
+        >>> A.columnwise_select(3, 6)
+        SparseMatrix(indices=tensor([[1, 2, 3],
+                                     [1, 0, 2]]),
+                     values=tensor([2, 3, 4]),
                      shape=(5, 3), nnz=3)
         """
-        if dim not in (0, 1):
-            raise ValueError("The selection dimension should be 0 or 1.")
-        if isinstance(index, slice):
-            raise NotImplementedError
-        raise TypeError(f"{type(index).__name__} is unsupported input type.")
+        return self.c_sparse_matrix.columnwise_select(start, end)
 
 
 def spmatrix(