Repository: pyg-team/pytorch_geometric
Pull request #7501 (2023-06-03 15:13:52)

Add `NodeEncoder` from the GraphMixer paper
Not very sure it's the correct way to implement what they did in GraphMixer, but hopefully not too far off..

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7501?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
> Merging [#7501](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7501?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (352df4b) into [master](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/93d5066257e7b366f297fa2aa8bdd1e781a6189a?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) (93d5066) will **increase** coverage by `1.84%`.
> The diff coverage is `100.00%`.

```diff
@@            Coverage Diff             @@
##           master    #7501      +/-   ##
==========================================
+ Coverage   89.58%   91.42%   +1.84%     
==========================================
  Files         434      446      +12     
  Lines       24463    24839     +376     
==========================================
+ Hits        21914    22709     +795     
+ Misses       2549     2130     -419     
```


| [Impacted Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7501?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/\_\_init\_\_.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7501?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL19faW5pdF9fLnB5) | `100.00% <100.00%> (ø)` | |
| [torch\_geometric/nn/encoding.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7501?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL2VuY29kaW5nLnB5) | `100.00% <100.00%> (ø)` | |

... and [45 files with indirect coverage changes](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/7501/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 f1c4f940d..7f22d25f8 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 the `GraphMixer` model ([#7501](https://github.com/pyg-team/pytorch_geometric/pull/7501))
 - Added the `disable_dynamic_shape` experimental flag ([#7246](https://github.com/pyg-team/pytorch_geometric/pull/7246))
 - Added the option to override `use_segmm` selection in `HeteroLinear` ([#7474](https://github.com/pyg-team/pytorch_geometric/pull/7474))
 - Added the `MovieLens-1M` heterogeneous dataset ([#7479](https://github.com/pyg-team/pytorch_geometric/pull/7479))
diff --git a/torch_geometric/nn/models/graph_mixer.py b/torch_geometric/nn/models/graph_mixer.py
new file mode 100644
index 000000000..dd702383a
--- /dev/null
+++ b/torch_geometric/nn/models/graph_mixer.py
@@ -0,0 +1,49 @@
+import torch
+from torch import Tensor
+
+from torch_geometric.utils import scatter
+
+
+class NodeEncoder(torch.nn.Module):
+    r"""The node encoder module from the `"Do We Really Need Complicated
+    Model Architectures for Temporal Networks?"
+    <https://openreview.net/forum?id=ayPPc0SyLv1>`_ paper.
+    :class:`NodeEncoder` captures the 1-hop temporal neighborhood information
+    via mean pooling.
+
+    .. math::
+        \mathbf{x}_v^{\prime}(t_0) = \mathbf{x}_v + \textrm{mean} \left\{
+        \mathbf{x}_w : w \in \mathcal{N}(v, t_0 - T, t_0) \right\}
+
+    Args:
+        time_window (int): The temporal window size :math:`T` to define the
+            1-hop temporal neighborhood.
+    """
+    def __init__(self, time_window: int):
+        super().__init__()
+        self.time_window = time_window
+
+    def forward(
+        self,
+        x: Tensor,
+        edge_index: Tensor,
+        edge_time: Tensor,
+        seed_time: Tensor,
+    ) -> Tensor:
+        r"""
+        Args:
+            x (torch.Tensor): The input node features.
+            edge_index (torch.Tensor): The edge indices.
+            edge_time (torch.Tensor): The timestamp attached to every edge.
+            seed_time (torch.Tensor): The seed time :math:`t_0` for every
+                destination node.
+        """
+        mask = ((edge_time <= seed_time[edge_index[1]]) &
+                (edge_time > seed_time[edge_index[1]] - self.time_window))
+
+        src, dst = edge_index[:, mask]
+        mean = scatter(x[src], dst, dim=0, dim_size=x.size(0), reduce='mean')
+        return x + mean
+
+    def __repr__(self) -> str:
+        return f'{self.__class__.__name__}(time_window={self.time_window})'


Test patch:
diff --git a/test/nn/models/test_graph_mixer.py b/test/nn/models/test_graph_mixer.py
new file mode 100644
index 000000000..072ad5061
--- /dev/null
+++ b/test/nn/models/test_graph_mixer.py
@@ -0,0 +1,26 @@
+import torch
+
+from torch_geometric.nn.models.graph_mixer import NodeEncoder
+
+
+def test_node_encoder():
+    x = torch.arange(4, dtype=torch.float).view(-1, 1)
+    edge_index = torch.tensor([[1, 2, 0, 0, 1, 3], [0, 0, 1, 2, 2, 2]])
+    edge_time = torch.tensor([0, 1, 1, 1, 2, 3])
+    seed_time = torch.tensor([2, 2, 2, 2])
+
+    encoder = NodeEncoder(time_window=2)
+    assert str(encoder) == 'NodeEncoder(time_window=2)'
+
+    out = encoder(x, edge_index, edge_time, seed_time)
+    # Node 0 aggregates information from node 2 (excluding node 1).
+    # Node 1 aggregates information from node 0.
+    # Node 2 aggregates information from node 0 and node 1 (exluding node 3).
+    # Node 3 aggregates no information.
+    expected = torch.tensor([
+        [0 + 2],
+        [1 + 0],
+        [2 + 0.5 * (0 + 1)],
+        [3],
+    ])
+    assert torch.allclose(out, expected)