Repository: pyg-team/pytorch_geometric
Pull request #8117 (2023-10-03 11:30:24)

Added `mask` argument to `dense_to_sparse`
Added optional argument mask to dense_to_sparse so that it can correctly invert a call to to_dense_adj by returning the correct edge_index in case there are graphs with different number of nodes (and hence, the dense adjacency matrix contains some padding)

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8117?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#8117](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8117?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (14e569c) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/4ac86965256898b2abc8f1f2c124bacc83df34d1?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (4ac8696) will **decrease** coverage by `0.01%`.
> The diff coverage is `91.66%`.

```diff
@@            Coverage Diff             @@
##           master    #8117      +/-   ##
==========================================
- Coverage   88.23%   88.23%   -0.01%     
==========================================
  Files         471      471              
  Lines       28227    28237      +10     
==========================================
+ Hits        24907    24915       +8     
- Misses       3320     3322       +2     
```


| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8117?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/utils/sparse.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8117?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL3V0aWxzL3NwYXJzZS5weQ==) | `90.00% <91.66%> (ø)` | |

... and [1 file with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8117/indirect-changes?src=pr&el=tree-more&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)

:mega: We’re building smart automated test selection to slash your CI/CD build times. [Learn more](https://about.codecov.io/iterative-testing/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team)


Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 1f9da5e5a..532378d31 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -7,6 +7,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 
 ### Added
 
+- Added support for a node-level `mask` tensor in `dense_to_sparse` ([#8117](https://github.com/pyg-team/pytorch_geometric/pull/8117))
 - Added the `to_on_disk_dataset()` method to convert `InMemoryDataset` instances to `OnDiskDataset` instances ([#8116](https://github.com/pyg-team/pytorch_geometric/pull/8116))
 - Added `torch-frame` support ([#8110](https://github.com/pyg-team/pytorch_geometric/pull/8110))
 - Added the `DistLoader` base class ([#8079](https://github.com/pyg-team/pytorch_geometric/pull/8079))
diff --git a/torch_geometric/utils/sparse.py b/torch_geometric/utils/sparse.py
index 47189a68d..f5b39b2ce 100644
--- a/torch_geometric/utils/sparse.py
+++ b/torch_geometric/utils/sparse.py
@@ -1,3 +1,4 @@
+import warnings
 from typing import Any, List, Optional, Tuple, Union
 
 import torch
@@ -5,23 +6,29 @@ from torch import Tensor
 
 import torch_geometric.typing
 from torch_geometric.typing import SparseTensor
-from torch_geometric.utils import coalesce
+from torch_geometric.utils import coalesce, cumsum
 
 
-def dense_to_sparse(adj: Tensor) -> Tuple[Tensor, Tensor]:
+def dense_to_sparse(
+    adj: Tensor,
+    mask: Optional[Tensor] = None,
+) -> Tuple[Tensor, Tensor]:
     r"""Converts a dense adjacency matrix to a sparse adjacency matrix defined
     by edge indices and edge attributes.
 
     Args:
-        adj (Tensor): The dense adjacency matrix of shape
+        adj (torch.Tensor): The dense adjacency matrix of shape
             :obj:`[num_nodes, num_nodes]` or
             :obj:`[batch_size, num_nodes, num_nodes]`.
+        mask (torch.Tensor, optional): A boolean tensor of shape
+            :obj:`[batch_size, num_nodes]` holding information about which
+            nodes are in each example are valid. (default: :obj:`None`)
 
     :rtype: (:class:`LongTensor`, :class:`Tensor`)
 
     Examples:
 
-        >>> # Forr a single adjacency matrix
+        >>> # For a single adjacency matrix:
         >>> adj = torch.tensor([[3, 1],
         ...                     [2, 0]])
         >>> dense_to_sparse(adj)
@@ -29,7 +36,7 @@ def dense_to_sparse(adj: Tensor) -> Tuple[Tensor, Tensor]:
                 [0, 1, 0]]),
         tensor([3, 1, 2]))
 
-        >>> # For two adjacency matrixes
+        >>> # For two adjacency matrixes:
         >>> adj = torch.tensor([[[3, 1],
         ...                      [2, 0]],
         ...                     [[0, 1],
@@ -38,21 +45,70 @@ def dense_to_sparse(adj: Tensor) -> Tuple[Tensor, Tensor]:
         (tensor([[0, 0, 1, 2, 3],
                 [0, 1, 0, 3, 3]]),
         tensor([3, 1, 2, 1, 2]))
+
+        >>> # First graph with two nodes, second with three:
+        >>> adj = torch.tensor([[
+        ...         [3, 1, 0],
+        ...         [2, 0, 0],
+        ...         [0, 0, 0]
+        ...     ], [
+        ...         [0, 1, 0],
+        ...         [0, 2, 3],
+        ...         [0, 5, 0]
+        ...     ]])
+        >>> mask = torch.tensor([
+        ...         [True, True, False],
+        ...         [True, True, True]
+        ...     ])
+        >>> dense_to_sparse(adj, mask)
+        (tensor([[0, 0, 1, 2, 3, 3, 4],
+                [0, 1, 0, 3, 3, 4, 3]]),
+        tensor([3, 1, 2, 1, 2, 3, 5]))
     """
     if adj.dim() < 2 or adj.dim() > 3:
-        raise ValueError(f"Dense adjacency matrix 'adj' must be 2- or "
-                         f"3-dimensional (got {adj.dim()} dimensions)")
+        raise ValueError(f"Dense adjacency matrix 'adj' must be two- or "
+                         f"three-dimensional (got {adj.dim()} dimensions)")
+
+    if mask is not None and adj.dim() == 2:
+        warnings.warn("Mask should not be provided in case the dense "
+                      "adjacency matrix is two-dimensional")
+        mask = None
 
-    edge_index = adj.nonzero().t()
+    if mask is not None and mask.dim() != 2:
+        raise ValueError(f"Mask must be two-dimensional "
+                         f"(got {mask.dim()} dimensions)")
 
-    if edge_index.size(0) == 2:
+    if mask is not None and adj.size(-2) != adj.size(-1):
+        raise ValueError(f"Mask is only supported on quadratic adjacency "
+                         f"matrices (got [*, {adj.size(-2)}, {adj.size(-1)}])")
+
+    if adj.dim() == 2:
+        edge_index = adj.nonzero().t()
         edge_attr = adj[edge_index[0], edge_index[1]]
         return edge_index, edge_attr
     else:
-        edge_attr = adj[edge_index[0], edge_index[1], edge_index[2]]
-        row = edge_index[1] + adj.size(-2) * edge_index[0]
-        col = edge_index[2] + adj.size(-1) * edge_index[0]
-        return torch.stack([row, col], dim=0), edge_attr
+        flatten_adj = adj.view(-1, adj.size(-1))
+        if mask is not None:
+            flatten_adj = flatten_adj[mask.view(-1)]
+        edge_index = flatten_adj.nonzero().t()
+        edge_attr = flatten_adj[edge_index[0], edge_index[1]]
+
+        if mask is None:
+            offset = torch.arange(
+                start=0,
+                end=adj.size(0) * adj.size(2),
+                step=adj.size(2),
+                device=adj.device,
+            )
+            offset = offset.repeat_interleave(adj.size(1))
+        else:
+            count = mask.sum(dim=-1)
+            offset = cumsum(count)[:-1]
+            offset = offset.repeat_interleave(count)
+
+        edge_index[1] += offset[edge_index[0]]
+
+        return edge_index, edge_attr
 
 
 def is_torch_sparse_tensor(src: Any) -> bool:


Test patch:
diff --git a/test/utils/test_sparse.py b/test/utils/test_sparse.py
index 6c1af5e8d..ba9ed6c5d 100644
--- a/test/utils/test_sparse.py
+++ b/test/utils/test_sparse.py
@@ -49,6 +49,33 @@ def test_dense_to_sparse():
         assert edge_index.tolist() == [[0, 0, 1, 2, 3], [0, 1, 0, 3, 3]]
         assert edge_attr.tolist() == [3, 1, 2, 1, 2]
 
+    adj = torch.tensor([
+        [
+            [3.0, 1.0, 0.0],
+            [2.0, 0.0, 0.0],
+            [0.0, 0.0, 0.0],
+        ],
+        [
+            [0.0, 1.0, 0.0],
+            [0.0, 2.0, 3.0],
+            [0.0, 5.0, 0.0],
+        ],
+    ])
+    mask = torch.tensor([[True, True, False], [True, True, True]])
+
+    edge_index, edge_attr = dense_to_sparse(adj, mask)
+
+    assert edge_index.tolist() == [[0, 0, 1, 2, 3, 3, 4],
+                                   [0, 1, 0, 3, 3, 4, 3]]
+    assert edge_attr.tolist() == [3, 1, 2, 1, 2, 3, 5]
+
+    if is_full_test():
+        jit = torch.jit.script(dense_to_sparse)
+        edge_index, edge_attr = jit(adj, mask)
+        assert edge_index.tolist() == [[0, 0, 1, 2, 3, 3, 4],
+                                       [0, 1, 0, 3, 3, 4, 3]]
+        assert edge_attr.tolist() == [3, 1, 2, 1, 2, 3, 5]
+
 
 def test_dense_to_sparse_bipartite():
     edge_index, edge_attr = dense_to_sparse(torch.rand(2, 10, 5))