fix name

dgl_sparse/include/sparse/matrix_ops.h

@@ -28,8 +28,8 @@ std::tuple<std::shared_ptr<COO>, torch::Tensor, torch::Tensor> COOIntersection(
     const std::shared_ptr<COO>& lhs, const std::shared_ptr<COO>& rhs);
 
 /**
- * @brief Relabels indices of a dimension and removes rows or columns without
- * non-zero elements in the sparse matrix.
+ * @brief Compact sparse matrix by removing rows or columns without non-zero
+ * elements in the sparse matrix and relabeling indices of the dimension.
  *
  * This function serves a dual purpose: it allows you to reorganize the
  * indices within a specific dimension (rows or columns) of the sparse matrix
@@ -45,7 +45,7 @@ std::tuple<std::shared_ptr<COO>, torch::Tensor, torch::Tensor> COOIntersection(
  * @return A tuple containing the relabeled sparse matrix and the index mapping
  *         of the relabeled dimension from the new index to the original index.
  */
-std::tuple<c10::intrusive_ptr<SparseMatrix>, torch::Tensor> Relabel(
+std::tuple<c10::intrusive_ptr<SparseMatrix>, torch::Tensor> Compact(
     const c10::intrusive_ptr<SparseMatrix>& mat, uint64_t dim,
     torch::Tensor leading_indices);
 

dgl_sparse/include/sparse/matrix_ops_impl.h

@@ -11,8 +11,26 @@
 namespace dgl {
 namespace sparse {
 
+/**
+ * @brief Compact sparse matrix by removing rows or columns without non-zero
+ * elements in the sparse matrix and relabeling indices of the dimension.
+ *
+ * This function serves a dual purpose: it allows you to reorganize the
+ * indices within a specific dimension (rows or columns) of the sparse matrix
+ * and, if needed, place certain 'leading_indices' at the beginning of the
+ * relabeled dimension.
+ *
+ * @param mat The sparse matrix to be relabeled.
+ * @param dim The dimension to relabel. Should be 0 or 1. Use 0 for row-wise
+ *        relabeling and 1 for column-wise relabeling.
+ * @param leading_indices An optional tensor containing row or column ids that
+ *        should be placed at the beginning of the relabeled dimension.
+ *
+ * @return A tuple containing the relabeled sparse matrix and the index mapping
+ *         of the relabeled dimension from the new index to the original index.
+ */
 template <c10::DeviceType XPU, typename IdType>
-std::tuple<c10::intrusive_ptr<SparseMatrix>, torch::Tensor> RelabelImpl(
+std::tuple<c10::intrusive_ptr<SparseMatrix>, torch::Tensor> CompactImpl(
     const c10::intrusive_ptr<SparseMatrix>& mat, uint64_t dim,
     torch::Tensor leading_indices) {}
 

dgl_sparse/src/matrix_ops.cc

@@ -59,8 +59,8 @@ std::tuple<std::shared_ptr<COO>, torch::Tensor, torch::Tensor> COOIntersection(
 }
 
 /**
- * @brief Relabels indices of a dimension and removes rows or columns without
- * non-zero elements in the sparse matrix.
+ * @brief Compact sparse matrix by removing rows or columns without non-zero
+ * elements in the sparse matrix and relabeling indices of the dimension.
  *
  * This function serves a dual purpose: it allows you to reorganize the
  * indices within a specific dimension (rows or columns) of the sparse matrix
@@ -76,11 +76,11 @@ std::tuple<std::shared_ptr<COO>, torch::Tensor, torch::Tensor> COOIntersection(
  * @return A tuple containing the relabeled sparse matrix and the index mapping
  *         of the relabeled dimension from the new index to the original index.
  */
-std::tuple<c10::intrusive_ptr<SparseMatrix>, torch::Tensor> Relabel(
+std::tuple<c10::intrusive_ptr<SparseMatrix>, torch::Tensor> Compact(
     const c10::intrusive_ptr<SparseMatrix>& mat, uint64_t dim,
     torch::Tensor leading_indices) {
-  DGL_SPARSE_COO_SWITCH(mat->COOPtr(), XPU, IdType, "Relabel", {
-    return RelabelImpl<XPU, IdType>(mat, dim, leading_indices);
+  DGL_SPARSE_COO_SWITCH(mat->COOPtr(), XPU, IdType, "Compact", {
+    return CompactImpl<XPU, IdType>(mat, dim, leading_indices);
   });
 }
 

python/dgl/sparse/sparse_matrix.py

@@ -680,13 +680,13 @@ def sample(
             self.c_sparse_matrix.sample(dim, fanout, ids, replace, bias)
         )
 
-    def relabel(
+    def compact(
         self,
         dim: int,
         leading_indices: Optional[torch.Tensor] = None,
     ):
-        """Relabels indices of a dimension and remove rows or columns without
-        non-zero elements in the sparse matrix.
+        """Compact sparse matrix by removing rows or columns without non-zero
+        elements in the sparse matrix and relabeling indices of the dimension.
 
         This function serves a dual purpose: it allows you to reorganize the
         indices within a specific dimension (rows or columns) of the sparse
@@ -697,7 +697,10 @@ def relabel(
         of relabeled indices remains the same as the original order, except that
         rows or columns without non-zero elements are removed. When
         'leading_indices' are provided, they are positioned at the start of the
-        relabeled dimension.
+        relabeled dimension. To be precise, all rows selected by the specified
+        indices will be remapped from 0 to length(indices) - 1. Rows that are not
+        selected and contain any non-zero elements will be positioned after those
+        remapped rows while maintaining their original order.
 
         This function mimics 'dgl.to_block', a method used to compress a sampled
         subgraph by eliminating redundant nodes. The 'leading_indices' parameter
@@ -727,22 +730,27 @@ def relabel(
         >>> indices = torch.tensor([[0, 2],
                                     [1, 2]])
         >>> A = dglsp.spmatrix(indices)
+        >>> print(A.to_dense())
+        tensor([[0., 1., 0.],
+                [0., 0., 0.],
+                [0., 0., 1.]])
 
-        Case 1: Relabel rows without indices.
+        Case 1: Compact rows without indices.
 
-        >>> B, original_rows = A.relabel(dim=0, leading_indices=None)
-        >>> print(B)
-        SparseMatrix(indices=tensor([[0, 1], [1, 2]]),
-                     shape=(2, 3), nnz=2)
+        >>> B, original_rows = A.compact(dim=0, leading_indices=None)
+        >>> print(B.to_dense())
+        tensor([[0., 1., 0.],
+                [0., 0., 1.]])
         >>> print(original_rows)
         torch.Tensor([0, 2])
 
-        Case 2: Relabel rows with indices.
+        Case 2: Compact rows with indices.
 
-        >>> B, original_rows = A.relabel(dim=0, leading_indices=[1, 2])
-        >>> print(B)
-        SparseMatrix(indices=tensor([[1, 2], [2, 1]]),
-                     shape=(3, 3), nnz=2)
+        >>> B, original_rows = A.compact(dim=0, leading_indices=[1, 2])
+        >>> print(B.to_dense())
+        tensor([[0., 0., 0.],
+                [0., 0., 1.],
+                [0., 1., 0.],])
         >>> print(original_rows)
         torch.Tensor([1, 2, 0])
         """