We are currently solving the following issue within our repository. Here is the issue text:
--- BEGIN 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_

--- END ISSUE ---

Below are some code segments, each from a relevant file. One or more of these files may contain bugs.
--- BEGIN FILES ---
Path: `torch_geometric/utils/sparse.py`
Content:
```
1 from typing import Any, Optional, Tuple, Union
2 
3 import torch
4 from torch import Tensor
5 
6 from torch_geometric.typing import SparseTensor
7 
8 
9 def dense_to_sparse(adj: Tensor) -> Tuple[Tensor, Tensor]:
10     r"""Converts a dense adjacency matrix to a sparse adjacency matrix defined
11     by edge indices and edge attributes.
12 
13     Args:
14         adj (Tensor): The dense adjacency matrix.
15     :rtype: (:class:`LongTensor`, :class:`Tensor`)
16 
17     Examples:
18 
19         >>> # Forr a single adjacency matrix
20         >>> adj = torch.tensor([[3, 1],
21         ...                     [2, 0]])
22         >>> dense_to_sparse(adj)
23         (tensor([[0, 0, 1],
24                 [0, 1, 0]]),
25         tensor([3, 1, 2]))
26 
27         >>> # For two adjacency matrixes
28         >>> adj = torch.tensor([[[3, 1],
29         ...                      [2, 0]],
30         ...                     [[0, 1],
31         ...                      [0, 2]]])
32         >>> dense_to_sparse(adj)
33         (tensor([[0, 0, 1, 2, 3],
34                 [0, 1, 0, 3, 3]]),
35         tensor([3, 1, 2, 1, 2]))
36     """
37     assert adj.dim() >= 2 and adj.dim() <= 3
38     assert adj.size(-1) == adj.size(-2)
39 
40     edge_index = adj.nonzero().t()
41 
42     if edge_index.size(0) == 2:
43         edge_attr = adj[edge_index[0], edge_index[1]]
44         return edge_index, edge_attr
45     else:
46         edge_attr = adj[edge_index[0], edge_index[1], edge_index[2]]
47         batch = edge_index[0] * adj.size(-1)
48         row = batch + edge_index[1]
49         col = batch + edge_index[2]
50         return torch.stack([row, col], dim=0), edge_attr
51 
52 
53 def is_torch_sparse_tensor(src: Any) -> bool:
54     """Returns :obj:`True` if the input :obj:`src` is a
55     :class:`torch.sparse.Tensor` (in any sparse layout).
56 
57     Args:
58         src (Any): The input object to be checked.
59     """
60     return isinstance(src, Tensor) and src.is_sparse
61 
62 
63 def is_sparse(src: Any) -> bool:
64     """Returns :obj:`True` if the input :obj:`src` is of type
65     :class:`torch.sparse.Tensor` (in any sparse layout) or of type
66     :class:`torch_sparse.SparseTensor`.
67 
68     Args:
69         src (Any): The input object to be checked.
70     """
71     return is_torch_sparse_tensor(src) or isinstance(src, SparseTensor)
72 
73 
74 def to_torch_coo_tensor(
75     edge_index: Tensor,
76     edge_attr: Optional[Tensor] = None,
77     size: Optional[Union[int, Tuple[int, int]]] = None,
78 ) -> Tensor:
79     """Converts a sparse adjacency matrix defined by edge indices and edge
80     attributes to a :class:`torch.sparse.Tensor`.
81 
82     Args:
83         edge_index (LongTensor): The edge indices.
84         edge_attr (Tensor, optional): The edge attributes.
85             (default: :obj:`None`)
86         size (int or (int, int), optional): The size of the sparse matrix.
87             If given as an integer, will create a quadratic sparse matrix.
88             If set to :obj:`None`, will infer a quadratic sparse matrix based
89             on :obj:`edge_index.max() + 1`. (default: :obj:`None`)
90 
91     :rtype: :class:`torch.sparse.FloatTensor`
92 
93     Example:
94 
95         >>> edge_index = torch.tensor([[0, 1, 1, 2, 2, 3],
96         ...                            [1, 0, 2, 1, 3, 2]])
97         >>> to_torch_coo_tensor(edge_index)
98         tensor(indices=tensor([[0, 1, 1, 2, 2, 3],
99                             [1, 0, 2, 1, 3, 2]]),
100             values=tensor([1., 1., 1., 1., 1., 1.]),
101             size=(4, 4), nnz=6, layout=torch.sparse_coo)
102 
103     """
104     if size is None:
105         size = int(edge_index.max()) + 1
106     if not isinstance(size, (tuple, list)):
107         size = (size, size)
108 
109     if edge_attr is None:
110         edge_attr = torch.ones(edge_index.size(1), device=edge_index.device)
111 
112     size = tuple(size) + edge_attr.size()[1:]
113     out = torch.sparse_coo_tensor(edge_index, edge_attr, size,
114                                   device=edge_index.device)
115     out = out.coalesce()
116     return out
117
```
--- END FILES ---

Please first localize the bug based on the issue statement, and then generate a patch according to the `git diff` format fenced by three backticks.

Here is an example:

```diff
diff --git a/examples/server_async.py b/examples/server_async.py
--- a/examples/server_async.py
+++ b/examples/server_async.py
@@ -313,4 +313,4 @@
 
 
 if __name__ == "__main__":
-    asyncio.run(run_async_server("."), debug=True)
+    asyncio.run(run_async_server(), debug=True)
diff --git a/examples/server_sync.py b/examples/server_sync.py
--- a/examples/server_sync.py
+++ b/examples/server_sync.py
@@ -313,5 +313,5 @@
 
 
 if __name__ == "__main__":
-    server = run_sync_server(".")
+    server = run_sync_server()
     server.shutdown()

```

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 @@
     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 @@
                 [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 @@
         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

{"golden_diff": "diff --git a/torch_geometric/utils/sparse.py b/torch_geometric/utils/sparse.py\n--- a/torch_geometric/utils/sparse.py\n+++ b/torch_geometric/utils/sparse.py\n@@ -11,7 +11,10 @@\n     by edge indices and edge attributes.\n \n     Args:\n-        adj (Tensor): The dense adjacency matrix.\n+        adj (Tensor): The dense adjacency matrix of shape\n+            :obj:`[num_nodes, num_nodes]` or\n+            :obj:`[batch_size, num_nodes, num_nodes]`.\n+\n     :rtype: (:class:`LongTensor`, :class:`Tensor`)\n \n     Examples:\n@@ -34,8 +37,9 @@\n                 [0, 1, 0, 3, 3]]),\n         tensor([3, 1, 2, 1, 2]))\n     \"\"\"\n-    assert adj.dim() >= 2 and adj.dim() <= 3\n-    assert adj.size(-1) == adj.size(-2)\n+    if adj.dim() < 2 or adj.dim() > 3:\n+        raise ValueError(f\"Dense adjacency matrix 'adj' must be 2- or \"\n+                         f\"3-dimensional (got {adj.dim()} dimensions)\")\n \n     edge_index = adj.nonzero().t()\n \n@@ -44,9 +48,8 @@\n         return edge_index, edge_attr\n     else:\n         edge_attr = adj[edge_index[0], edge_index[1], edge_index[2]]\n-        batch = edge_index[0] * adj.size(-1)\n-        row = batch + edge_index[1]\n-        col = batch + edge_index[2]\n+        row = edge_index[1] + adj.size(-2) * edge_index[0]\n+        col = edge_index[2] + adj.size(-1) * edge_index[0]\n         return torch.stack([row, col], dim=0), edge_attr\n", "issue": "Bipartite graph support for utils.dense_to_sparse\n### \ud83d\ude80 The feature, motivation and pitch\n\nI 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).\n\n### Alternatives\n\nA 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.\n\n### Additional context\n\n_No response_\n", "before_files": [{"content": "from typing import Any, Optional, Tuple, Union\n\nimport torch\nfrom torch import Tensor\n\nfrom torch_geometric.typing import SparseTensor\n\n\ndef dense_to_sparse(adj: Tensor) -> Tuple[Tensor, Tensor]:\n    r\"\"\"Converts a dense adjacency matrix to a sparse adjacency matrix defined\n    by edge indices and edge attributes.\n\n    Args:\n        adj (Tensor): The dense adjacency matrix.\n    :rtype: (:class:`LongTensor`, :class:`Tensor`)\n\n    Examples:\n\n        >>> # Forr a single adjacency matrix\n        >>> adj = torch.tensor([[3, 1],\n        ...                     [2, 0]])\n        >>> dense_to_sparse(adj)\n        (tensor([[0, 0, 1],\n                [0, 1, 0]]),\n        tensor([3, 1, 2]))\n\n        >>> # For two adjacency matrixes\n        >>> adj = torch.tensor([[[3, 1],\n        ...                      [2, 0]],\n        ...                     [[0, 1],\n        ...                      [0, 2]]])\n        >>> dense_to_sparse(adj)\n        (tensor([[0, 0, 1, 2, 3],\n                [0, 1, 0, 3, 3]]),\n        tensor([3, 1, 2, 1, 2]))\n    \"\"\"\n    assert adj.dim() >= 2 and adj.dim() <= 3\n    assert adj.size(-1) == adj.size(-2)\n\n    edge_index = adj.nonzero().t()\n\n    if edge_index.size(0) == 2:\n        edge_attr = adj[edge_index[0], edge_index[1]]\n        return edge_index, edge_attr\n    else:\n        edge_attr = adj[edge_index[0], edge_index[1], edge_index[2]]\n        batch = edge_index[0] * adj.size(-1)\n        row = batch + edge_index[1]\n        col = batch + edge_index[2]\n        return torch.stack([row, col], dim=0), edge_attr\n\n\ndef is_torch_sparse_tensor(src: Any) -> bool:\n    \"\"\"Returns :obj:`True` if the input :obj:`src` is a\n    :class:`torch.sparse.Tensor` (in any sparse layout).\n\n    Args:\n        src (Any): The input object to be checked.\n    \"\"\"\n    return isinstance(src, Tensor) and src.is_sparse\n\n\ndef is_sparse(src: Any) -> bool:\n    \"\"\"Returns :obj:`True` if the input :obj:`src` is of type\n    :class:`torch.sparse.Tensor` (in any sparse layout) or of type\n    :class:`torch_sparse.SparseTensor`.\n\n    Args:\n        src (Any): The input object to be checked.\n    \"\"\"\n    return is_torch_sparse_tensor(src) or isinstance(src, SparseTensor)\n\n\ndef to_torch_coo_tensor(\n    edge_index: Tensor,\n    edge_attr: Optional[Tensor] = None,\n    size: Optional[Union[int, Tuple[int, int]]] = None,\n) -> Tensor:\n    \"\"\"Converts a sparse adjacency matrix defined by edge indices and edge\n    attributes to a :class:`torch.sparse.Tensor`.\n\n    Args:\n        edge_index (LongTensor): The edge indices.\n        edge_attr (Tensor, optional): The edge attributes.\n            (default: :obj:`None`)\n        size (int or (int, int), optional): The size of the sparse matrix.\n            If given as an integer, will create a quadratic sparse matrix.\n            If set to :obj:`None`, will infer a quadratic sparse matrix based\n            on :obj:`edge_index.max() + 1`. (default: :obj:`None`)\n\n    :rtype: :class:`torch.sparse.FloatTensor`\n\n    Example:\n\n        >>> edge_index = torch.tensor([[0, 1, 1, 2, 2, 3],\n        ...                            [1, 0, 2, 1, 3, 2]])\n        >>> to_torch_coo_tensor(edge_index)\n        tensor(indices=tensor([[0, 1, 1, 2, 2, 3],\n                            [1, 0, 2, 1, 3, 2]]),\n            values=tensor([1., 1., 1., 1., 1., 1.]),\n            size=(4, 4), nnz=6, layout=torch.sparse_coo)\n\n    \"\"\"\n    if size is None:\n        size = int(edge_index.max()) + 1\n    if not isinstance(size, (tuple, list)):\n        size = (size, size)\n\n    if edge_attr is None:\n        edge_attr = torch.ones(edge_index.size(1), device=edge_index.device)\n\n    size = tuple(size) + edge_attr.size()[1:]\n    out = torch.sparse_coo_tensor(edge_index, edge_attr, size,\n                                  device=edge_index.device)\n    out = out.coalesce()\n    return out\n", "path": "torch_geometric/utils/sparse.py"}], "after_files": [{"content": "from typing import Any, Optional, Tuple, Union\n\nimport torch\nfrom torch import Tensor\n\nfrom torch_geometric.typing import SparseTensor\n\n\ndef dense_to_sparse(adj: Tensor) -> Tuple[Tensor, Tensor]:\n    r\"\"\"Converts a dense adjacency matrix to a sparse adjacency matrix defined\n    by edge indices and edge attributes.\n\n    Args:\n        adj (Tensor): The dense adjacency matrix of shape\n            :obj:`[num_nodes, num_nodes]` or\n            :obj:`[batch_size, num_nodes, num_nodes]`.\n\n    :rtype: (:class:`LongTensor`, :class:`Tensor`)\n\n    Examples:\n\n        >>> # Forr a single adjacency matrix\n        >>> adj = torch.tensor([[3, 1],\n        ...                     [2, 0]])\n        >>> dense_to_sparse(adj)\n        (tensor([[0, 0, 1],\n                [0, 1, 0]]),\n        tensor([3, 1, 2]))\n\n        >>> # For two adjacency matrixes\n        >>> adj = torch.tensor([[[3, 1],\n        ...                      [2, 0]],\n        ...                     [[0, 1],\n        ...                      [0, 2]]])\n        >>> dense_to_sparse(adj)\n        (tensor([[0, 0, 1, 2, 3],\n                [0, 1, 0, 3, 3]]),\n        tensor([3, 1, 2, 1, 2]))\n    \"\"\"\n    if adj.dim() < 2 or adj.dim() > 3:\n        raise ValueError(f\"Dense adjacency matrix 'adj' must be 2- or \"\n                         f\"3-dimensional (got {adj.dim()} dimensions)\")\n\n    edge_index = adj.nonzero().t()\n\n    if edge_index.size(0) == 2:\n        edge_attr = adj[edge_index[0], edge_index[1]]\n        return edge_index, edge_attr\n    else:\n        edge_attr = adj[edge_index[0], edge_index[1], edge_index[2]]\n        row = edge_index[1] + adj.size(-2) * edge_index[0]\n        col = edge_index[2] + adj.size(-1) * edge_index[0]\n        return torch.stack([row, col], dim=0), edge_attr\n\n\ndef is_torch_sparse_tensor(src: Any) -> bool:\n    \"\"\"Returns :obj:`True` if the input :obj:`src` is a\n    :class:`torch.sparse.Tensor` (in any sparse layout).\n\n    Args:\n        src (Any): The input object to be checked.\n    \"\"\"\n    return isinstance(src, Tensor) and src.is_sparse\n\n\ndef is_sparse(src: Any) -> bool:\n    \"\"\"Returns :obj:`True` if the input :obj:`src` is of type\n    :class:`torch.sparse.Tensor` (in any sparse layout) or of type\n    :class:`torch_sparse.SparseTensor`.\n\n    Args:\n        src (Any): The input object to be checked.\n    \"\"\"\n    return is_torch_sparse_tensor(src) or isinstance(src, SparseTensor)\n\n\ndef to_torch_coo_tensor(\n    edge_index: Tensor,\n    edge_attr: Optional[Tensor] = None,\n    size: Optional[Union[int, Tuple[int, int]]] = None,\n) -> Tensor:\n    \"\"\"Converts a sparse adjacency matrix defined by edge indices and edge\n    attributes to a :class:`torch.sparse.Tensor`.\n\n    Args:\n        edge_index (LongTensor): The edge indices.\n        edge_attr (Tensor, optional): The edge attributes.\n            (default: :obj:`None`)\n        size (int or (int, int), optional): The size of the sparse matrix.\n            If given as an integer, will create a quadratic sparse matrix.\n            If set to :obj:`None`, will infer a quadratic sparse matrix based\n            on :obj:`edge_index.max() + 1`. (default: :obj:`None`)\n\n    :rtype: :class:`torch.sparse.FloatTensor`\n\n    Example:\n\n        >>> edge_index = torch.tensor([[0, 1, 1, 2, 2, 3],\n        ...                            [1, 0, 2, 1, 3, 2]])\n        >>> to_torch_coo_tensor(edge_index)\n        tensor(indices=tensor([[0, 1, 1, 2, 2, 3],\n                            [1, 0, 2, 1, 3, 2]]),\n            values=tensor([1., 1., 1., 1., 1., 1.]),\n            size=(4, 4), nnz=6, layout=torch.sparse_coo)\n\n    \"\"\"\n    if size is None:\n        size = int(edge_index.max()) + 1\n    if not isinstance(size, (tuple, list)):\n        size = (size, size)\n\n    if edge_attr is None:\n        edge_attr = torch.ones(edge_index.size(1), device=edge_index.device)\n\n    size = tuple(size) + edge_attr.size()[1:]\n    out = torch.sparse_coo_tensor(edge_index, edge_attr, size,\n                                  device=edge_index.device)\n    out = out.coalesce()\n    return out\n", "path": "torch_geometric/utils/sparse.py"}]}