Added mask argument to dense_to_sparse

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))

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))

torch_geometric/utils/sparse.py

@@ -1,35 +1,42 @@
+import warnings
 from typing import Any, List, Optional, Tuple, Union
 
 import torch
 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)
         (tensor([[0, 0, 1],
                 [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: