Repository: pyg-team/pytorch_geometric
Pull request #8048 (2023-09-18 00:33:01)

Add `edge_attr` support for `ResGatedGraphConv`

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8048?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#8048](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8048?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (ba32fc9) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/d2f9d18eb1320b50964efe20b80ea107e35cfffd?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (d2f9d18) will **decrease** coverage by `0.01%`.
> Report is 3 commits behind head on master.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##           master    #8048      +/-   ##
==========================================
- Coverage   88.59%   88.59%   -0.01%     
==========================================
  Files         466      466              
  Lines       27679    27692      +13     
==========================================
+ Hits        24523    24534      +11     
- Misses       3156     3158       +2     
```


| [Files Changed](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8048?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Coverage Δ | |
|---|---|---|
| [torch\_geometric/nn/conv/res\_gated\_graph\_conv.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8048?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2NvbnYvcmVzX2dhdGVkX2dyYXBoX2NvbnYucHk=) | `100.00% <100.00%> (ø)` | |

... and [3 files with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/8048/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 98e0776c0..cad4b1758 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 `edge_attr` support to `ResGatedGraphConv` ([#8048](https://github.com/pyg-team/pytorch_geometric/pull/8048))
 - Added a `Database` interface and `SQLiteDatabase`/`RocksDatabase` implementations ([#8028](https://github.com/pyg-team/pytorch_geometric/pull/8028), [#8044](https://github.com/pyg-team/pytorch_geometric/pull/8044), [#8046](https://github.com/pyg-team/pytorch_geometric/pull/8046), [#8051](https://github.com/pyg-team/pytorch_geometric/pull/8051), [#8052](https://github.com/pyg-team/pytorch_geometric/pull/8052), [#8054](https://github.com/pyg-team/pytorch_geometric/pull/8054), [#8057](https://github.com/pyg-team/pytorch_geometric/pull/8057))
 - Added support for weighted/biased sampling in `NeighborLoader`/`LinkNeighborLoader` ([#8038](https://github.com/pyg-team/pytorch_geometric/pull/8038))
 - Added the `MixHopConv` layer and an corresponding example ([#8025](https://github.com/pyg-team/pytorch_geometric/pull/8025))
diff --git a/torch_geometric/nn/conv/res_gated_graph_conv.py b/torch_geometric/nn/conv/res_gated_graph_conv.py
index f9bd24128..7f9a4ef97 100644
--- a/torch_geometric/nn/conv/res_gated_graph_conv.py
+++ b/torch_geometric/nn/conv/res_gated_graph_conv.py
@@ -1,12 +1,13 @@
 from typing import Callable, Optional, Tuple, Union
 
+import torch
 from torch import Tensor
 from torch.nn import Parameter, Sigmoid
 
 from torch_geometric.nn.conv import MessagePassing
 from torch_geometric.nn.dense.linear import Linear
 from torch_geometric.nn.inits import zeros
-from torch_geometric.typing import Adj, PairTensor
+from torch_geometric.typing import Adj, OptTensor, PairTensor
 
 
 class ResGatedGraphConv(MessagePassing):
@@ -33,6 +34,8 @@ class ResGatedGraphConv(MessagePassing):
         out_channels (int): Size of each output sample.
         act (callable, optional): Gating function :math:`\sigma`.
             (default: :meth:`torch.nn.Sigmoid()`)
+        edge_dim (int, optional): Edge feature dimensionality (in case
+            there are any). (default: :obj:`None`)
         bias (bool, optional): If set to :obj:`False`, the layer will not learn
             an additive bias. (default: :obj:`True`)
         root_weight (bool, optional): If set to :obj:`False`, the layer will
@@ -55,6 +58,7 @@ class ResGatedGraphConv(MessagePassing):
         in_channels: Union[int, Tuple[int, int]],
         out_channels: int,
         act: Optional[Callable] = Sigmoid(),
+        edge_dim: Optional[int] = None,
         root_weight: bool = True,
         bias: bool = True,
         **kwargs,
@@ -66,14 +70,16 @@ class ResGatedGraphConv(MessagePassing):
         self.in_channels = in_channels
         self.out_channels = out_channels
         self.act = act
+        self.edge_dim = edge_dim
         self.root_weight = root_weight
 
         if isinstance(in_channels, int):
             in_channels = (in_channels, in_channels)
 
-        self.lin_key = Linear(in_channels[1], out_channels)
-        self.lin_query = Linear(in_channels[0], out_channels)
-        self.lin_value = Linear(in_channels[0], out_channels)
+        edge_dim = edge_dim if edge_dim is not None else 0
+        self.lin_key = Linear(in_channels[1] + edge_dim, out_channels)
+        self.lin_query = Linear(in_channels[0] + edge_dim, out_channels)
+        self.lin_value = Linear(in_channels[0] + edge_dim, out_channels)
 
         if root_weight:
             self.lin_skip = Linear(in_channels[1], out_channels, bias=False)
@@ -97,16 +103,24 @@ class ResGatedGraphConv(MessagePassing):
         if self.bias is not None:
             zeros(self.bias)
 
-    def forward(self, x: Union[Tensor, PairTensor], edge_index: Adj) -> Tensor:
+    def forward(self, x: Union[Tensor, PairTensor], edge_index: Adj,
+                edge_attr: OptTensor = None) -> Tensor:
+
         if isinstance(x, Tensor):
             x: PairTensor = (x, x)
 
-        k = self.lin_key(x[1])
-        q = self.lin_query(x[0])
-        v = self.lin_value(x[0])
+        # In case edge features are not given, we can compute key, query and
+        # value tensors in node-level space, which is a bit more efficient:
+        if self.edge_dim is None:
+            k = self.lin_key(x[1])
+            q = self.lin_query(x[0])
+            v = self.lin_value(x[0])
+        else:
+            k, q, v = x[1], x[0], x[0]
 
-        # propagate_type: (k: Tensor, q: Tensor, v: Tensor)
-        out = self.propagate(edge_index, k=k, q=q, v=v, size=None)
+        # propagate_type: (k: Tensor, q: Tensor, v: Tensor, edge_attr: OptTensor)  # noqa
+        out = self.propagate(edge_index, k=k, q=q, v=v, edge_attr=edge_attr,
+                             size=None)
 
         if self.root_weight:
             out = out + self.lin_skip(x[1])
@@ -116,5 +130,14 @@ class ResGatedGraphConv(MessagePassing):
 
         return out
 
-    def message(self, k_i: Tensor, q_j: Tensor, v_j: Tensor) -> Tensor:
+    def message(self, k_i: Tensor, q_j: Tensor, v_j: Tensor,
+                edge_attr: OptTensor) -> Tensor:
+
+        assert (edge_attr is not None) == (self.edge_dim is not None)
+
+        if edge_attr is not None:
+            k_i = self.lin_key(torch.cat([k_i, edge_attr], dim=-1))
+            q_j = self.lin_query(torch.cat([q_j, edge_attr], dim=-1))
+            v_j = self.lin_value(torch.cat([v_j, edge_attr], dim=-1))
+
         return self.act(k_i + q_j) * v_j


Test patch:
diff --git a/test/nn/conv/test_res_gated_graph_conv.py b/test/nn/conv/test_res_gated_graph_conv.py
index 193cd5381..8bebaa018 100644
--- a/test/nn/conv/test_res_gated_graph_conv.py
+++ b/test/nn/conv/test_res_gated_graph_conv.py
@@ -1,3 +1,4 @@
+import pytest
 import torch
 
 import torch_geometric.typing
@@ -7,53 +8,56 @@ from torch_geometric.typing import SparseTensor
 from torch_geometric.utils import to_torch_csc_tensor
 
 
-def test_res_gated_graph_conv():
+@pytest.mark.parametrize('edge_dim', [None, 4])
+def test_res_gated_graph_conv(edge_dim):
     x1 = torch.randn(4, 8)
     x2 = torch.randn(2, 32)
     edge_index = torch.tensor([[0, 1, 2, 3], [0, 0, 1, 1]])
+    edge_attr = torch.randn(edge_index.size(1), edge_dim) if edge_dim else None
     adj1 = to_torch_csc_tensor(edge_index, size=(4, 4))
 
-    conv = ResGatedGraphConv(8, 32)
+    conv = ResGatedGraphConv(8, 32, edge_dim=edge_dim)
     assert str(conv) == 'ResGatedGraphConv(8, 32)'
 
-    out = conv(x1, edge_index)
+    out = conv(x1, edge_index, edge_attr)
     assert out.size() == (4, 32)
-    assert torch.allclose(conv(x1, adj1.t()), out, atol=1e-6)
+    assert torch.allclose(conv(x1, adj1.t(), edge_attr), out, atol=1e-6)
 
     if torch_geometric.typing.WITH_TORCH_SPARSE:
-        adj2 = SparseTensor.from_edge_index(edge_index, sparse_sizes=(4, 4))
+        adj2 = SparseTensor.from_edge_index(edge_index, edge_attr, (4, 4))
         assert torch.allclose(conv(x1, adj2.t()), out, atol=1e-6)
 
     if is_full_test():
-        t = '(Tensor, Tensor) -> Tensor'
+        t = '(Tensor, Tensor, OptTensor) -> Tensor'
         jit = torch.jit.script(conv.jittable(t))
-        assert torch.allclose(jit(x1, edge_index), out, atol=1e-6)
+        assert torch.allclose(jit(x1, edge_index, edge_attr), out, atol=1e-6)
 
     if is_full_test() and torch_geometric.typing.WITH_TORCH_SPARSE:
-        t = '(Tensor, SparseTensor) -> Tensor'
+        t = '(Tensor, SparseTensor, OptTensor) -> Tensor'
         jit = torch.jit.script(conv.jittable(t))
         assert torch.allclose(jit(x1, adj2.t()), out, atol=1e-6)
 
     # Test bipartite message passing:
     adj1 = to_torch_csc_tensor(edge_index, size=(4, 2))
 
-    conv = ResGatedGraphConv((8, 32), 32)
+    conv = ResGatedGraphConv((8, 32), 32, edge_dim=edge_dim)
     assert str(conv) == 'ResGatedGraphConv((8, 32), 32)'
 
-    out = conv((x1, x2), edge_index)
+    out = conv((x1, x2), edge_index, edge_attr)
     assert out.size() == (2, 32)
-    assert torch.allclose(conv((x1, x2), adj1.t()), out, atol=1e-6)
+    assert torch.allclose(conv((x1, x2), adj1.t(), edge_attr), out, atol=1e-6)
 
     if torch_geometric.typing.WITH_TORCH_SPARSE:
-        adj2 = SparseTensor.from_edge_index(edge_index, sparse_sizes=(4, 2))
+        adj2 = SparseTensor.from_edge_index(edge_index, edge_attr, (4, 2))
         assert torch.allclose(conv((x1, x2), adj2.t()), out, atol=1e-6)
 
     if is_full_test():
-        t = '(PairTensor, Tensor) -> Tensor'
+        t = '(PairTensor, Tensor, OptTensor) -> Tensor'
         jit = torch.jit.script(conv.jittable(t))
-        assert torch.allclose(jit((x1, x2), edge_index), out, atol=1e-6)
+        assert torch.allclose(jit((x1, x2), edge_index, edge_attr), out,
+                              atol=1e-6)
 
     if is_full_test() and torch_geometric.typing.WITH_TORCH_SPARSE:
-        t = '(PairTensor, SparseTensor) -> Tensor'
+        t = '(PairTensor, SparseTensor, OptTensor) -> Tensor'
         jit = torch.jit.script(conv.jittable(t))
         assert torch.allclose(jit((x1, x2), adj2.t()), out, atol=1e-6)