diff --git a/python/dgl/sparse/sparse_matrix.py b/python/dgl/sparse/sparse_matrix.py index 9461a31eead0..8f3d60b46356 100644 --- a/python/dgl/sparse/sparse_matrix.py +++ b/python/dgl/sparse/sparse_matrix.py @@ -678,28 +678,33 @@ def relabel( 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. - 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 'heading_indices' at the beginning of the relabeled dimension. - - In the absence of 'heading_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 'heading_indices' are provided, - they are positioned at the start of the relabeled dimension. - - This function mimics 'dgl.to_block', a method used to compress a sampled - subgraph by eliminating redundant nodes. The 'heading_indices' parameter - replicates the behavior of 'include_dst_in_src' in 'dgl.to_block.' Setting - 'heading_indices' to column IDs when relabeling the row dimension, for example, - achieves the same effect as including destination nodes in source nodes. + """Relabels indices of a dimension and remove rows or columns without + non-zero elements in the sparse matrix. + + 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. + + 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.' + 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. - heading_indices : torch.Tensor, optional + leading_indices : torch.Tensor, optional An optional tensor containing row or column ids that should be placed at the beginning of the relabeled dimension. @@ -717,7 +722,7 @@ def relabel( Case 1: Relabel rows without indices. - >>> B, original_rows = A.relabel(dim=0, heading_indices=None) + >>> B, original_rows = A.relabel(dim=0, leading_indices=None) >>> print(B) SparseMatrix(indices=tensor([[0, 1], [1, 2]]), shape=(2, 3), nnz=2) @@ -726,7 +731,7 @@ def relabel( Case 2: Relabel rows with indices. - >>> B, original_rows = A.relabel(dim=0, heading_indices=[1, 2]) + >>> 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)