[Sparse] Add relabel python API (#6323)

python/dgl/sparse/sparse_matrix.py

@@ -680,6 +680,82 @@ class SparseMatrix:
             self.c_sparse_matrix.sample(dim, fanout, ids, replace, bias)
         )
 
+    def compact(
+        self,
+        dim: int,
+        leading_indices: Optional[torch.Tensor] = None,
+    ):
+        """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.
+
+        In the absence of 'leading_indices' (when it's set to `None`), the order
+        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. 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
+        replicates the behavior of 'include_dst_in_src' in 'dgl.to_block',
+        adding destination node information for message passing.
+        Setting 'leading_indices' to column IDs when relabeling the row
+        dimension, for example, achieves the same effect as including destination
+        nodes in source nodes.
+
+        Parameters
+        ----------
+        dim : int
+            The dimension to relabel. Should be 0 or 1. Use `dim = 0` for rowwise
+            relabeling and `dim = 1` for columnwise relabeling.
+        leading_indices : torch.Tensor, optional
+            An optional tensor containing row or column ids that should be placed
+            at the beginning of the relabeled dimension.
+
+        Returns
+        -------
+        Tuple[SparseMatrix, torch.Tensor]
+            A tuple containing the relabeled sparse matrix and the index mapping
+            of the relabeled dimension from the new index to the original index.
+
+        Examples
+        --------
+        >>> 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: Compact rows without indices.
+
+        >>> 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: Compact rows with indices.
+
+        >>> 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])
+        """
+        raise NotImplementedError
+
 
 def spmatrix(
     indices: torch.Tensor,