Repository: pyg-team/pytorch_geometric

Issue:
Bipartite graph support for utils.dense_to_sparse
### 🚀 The feature, motivation and pitch

I have a nearly-dense bipartite graph (that is, most features in node set A are connected to most features in node set B), and so it is easiest for me to define the edge adjacency matrix as a dense, non-square matrix. However, the message passing class expects a sparse edge adjacency layout. The dense_to_sparse utility would seem to be ideal for this purpose, but it can only take square matrices (thus, is unhelpful for bipartite graphs).

### Alternatives

A way to implicitly request propagate to pass messages from every node in A to every node in B would be even better (storing fully connected graphs is very memory inefficient), but I know that pyg is meant for sparser graph constructions so this would likely be a feature that wasn't used very much by other people.

### Additional context

_No response_


Discussion:
Thanks for this feature request, we should definitely work on this. Do you have interest in helping with this?
I'm not at all familiar with the back-end and code base for pyg yet, just using it and the documentation. I can try to help, though.

Resolving patch:
diff --git a/torch_geometric/utils/sparse.py b/torch_geometric/utils/sparse.py
--- a/torch_geometric/utils/sparse.py
+++ b/torch_geometric/utils/sparse.py
@@ -11,7 +11,10 @@ def dense_to_sparse(adj: Tensor) -> Tuple[Tensor, Tensor]:
     by edge indices and edge attributes.
 
     Args:
-        adj (Tensor): The dense adjacency matrix.
+        adj (Tensor): The dense adjacency matrix of shape
+            :obj:`[num_nodes, num_nodes]` or
+            :obj:`[batch_size, num_nodes, num_nodes]`.
+
     :rtype: (:class:`LongTensor`, :class:`Tensor`)
 
     Examples:
@@ -34,8 +37,9 @@ def dense_to_sparse(adj: Tensor) -> Tuple[Tensor, Tensor]:
                 [0, 1, 0, 3, 3]]),
         tensor([3, 1, 2, 1, 2]))
     """
-    assert adj.dim() >= 2 and adj.dim() <= 3
-    assert adj.size(-1) == adj.size(-2)
+    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)")
 
     edge_index = adj.nonzero().t()
 
@@ -44,9 +48,8 @@ def dense_to_sparse(adj: Tensor) -> Tuple[Tensor, Tensor]:
         return edge_index, edge_attr
     else:
         edge_attr = adj[edge_index[0], edge_index[1], edge_index[2]]
-        batch = edge_index[0] * adj.size(-1)
-        row = batch + edge_index[1]
-        col = batch + 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
 
 


Test patch:
diff --git a/test/utils/test_sparse.py b/test/utils/test_sparse.py
--- a/test/utils/test_sparse.py
+++ b/test/utils/test_sparse.py
@@ -43,6 +43,12 @@ def test_dense_to_sparse():
         assert edge_attr.tolist() == [3, 1, 2, 1, 2]
 
 
+def test_dense_to_sparse_bipartite():
+    edge_index, edge_attr = dense_to_sparse(torch.rand(2, 10, 5))
+    assert edge_index[0].max() == 19
+    assert edge_index[1].max() == 9
+
+
 def test_is_torch_sparse_tensor():
     x = torch.randn(5, 5)