Repository: pyg-team/pytorch_geometric
Pull request #9380 (2024-05-31 19:36:46)

add `HeteroJK` class
Hi,

this PR implements #9355, please do let me know if any additional changes or documentation is needed. 
Thank you

codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9380?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
Attention: Patch coverage is `89.47368%` with `2 lines` in your changes missing coverage. Please review.
> Project coverage is 87.86%. Comparing base [(`476e768`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/476e768593d91c1622b02440e25bbc93fc7f5294?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) to head [(`3866ae5`)](https://app.codecov.io/gh/pyg-team/pytorch_geometric/commit/3866ae50c5879eebf268e8a75e9e750d8d3cacc6?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).

> :exclamation: **Current head 3866ae5 differs from pull request most recent head 4d5ae36**
> 
> Please [upload](https://docs.codecov.com/docs/codecov-uploader) reports for the commit 4d5ae36 to get more accurate results.

| [Files](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9380?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Patch % | Lines |
|---|---|---|
| [torch\_geometric/nn/models/jumping\_knowledge.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9380?src=pr&el=tree&filepath=torch_geometric%2Fnn%2Fmodels%2Fjumping_knowledge.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL25uL21vZGVscy9qdW1waW5nX2tub3dsZWRnZS5weQ==) | 88.88% | [2 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9380?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) |

<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #9380      +/-   ##
==========================================
- Coverage   87.99%   87.86%   -0.13%     
==========================================
  Files         473      473              
  Lines       30820    30830      +10     
==========================================
- Hits        27120    27090      -30     
- Misses       3700     3740      +40     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/9380?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?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 1c6bd4979..2b7576128 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 heterogeneous `HeteroJumpingKnowledge` module ([#9380](https://github.com/pyg-team/pytorch_geometric/pull/9380))
 - Started work on GNN+LLM package ([#9350](https://github.com/pyg-team/pytorch_geometric/pull/9350))
 - Added support for negative sampling in `LinkLoader` acccording to source and destination node weights ([#9316](https://github.com/pyg-team/pytorch_geometric/pull/9316))
 - Added support for `EdgeIndex.unbind` ([#9298](https://github.com/pyg-team/pytorch_geometric/pull/9298))
diff --git a/torch_geometric/nn/models/__init__.py b/torch_geometric/nn/models/__init__.py
index 5494311e9..334970da5 100644
--- a/torch_geometric/nn/models/__init__.py
+++ b/torch_geometric/nn/models/__init__.py
@@ -2,7 +2,7 @@ r"""Model package."""
 
 from .mlp import MLP
 from .basic_gnn import GCN, GraphSAGE, GIN, GAT, PNA, EdgeCNN
-from .jumping_knowledge import JumpingKnowledge
+from .jumping_knowledge import JumpingKnowledge, HeteroJumpingKnowledge
 from .meta import MetaLayer
 from .node2vec import Node2Vec
 from .deep_graph_infomax import DeepGraphInfomax
@@ -42,6 +42,7 @@ __all__ = classes = [
     'PNA',
     'EdgeCNN',
     'JumpingKnowledge',
+    'HeteroJumpingKnowledge',
     'MetaLayer',
     'Node2Vec',
     'DeepGraphInfomax',
diff --git a/torch_geometric/nn/models/jumping_knowledge.py b/torch_geometric/nn/models/jumping_knowledge.py
index 02f79a0f3..3350d5617 100644
--- a/torch_geometric/nn/models/jumping_knowledge.py
+++ b/torch_geometric/nn/models/jumping_knowledge.py
@@ -1,4 +1,4 @@
-from typing import List, Optional
+from typing import Dict, List, Optional
 
 import torch
 from torch import Tensor
@@ -41,8 +41,12 @@ class JumpingKnowledge(torch.nn.Module):
         num_layers (int, optional): The number of layers to aggregate. Needs to
             be only set for LSTM-style aggregation. (default: :obj:`None`)
     """
-    def __init__(self, mode: str, channels: Optional[int] = None,
-                 num_layers: Optional[int] = None):
+    def __init__(
+        self,
+        mode: str,
+        channels: Optional[int] = None,
+        num_layers: Optional[int] = None,
+    ) -> None:
         super().__init__()
         self.mode = mode.lower()
         assert self.mode in ['cat', 'max', 'lstm']
@@ -63,7 +67,7 @@ class JumpingKnowledge(torch.nn.Module):
 
         self.reset_parameters()
 
-    def reset_parameters(self):
+    def reset_parameters(self) -> None:
         r"""Resets all learnable parameters of the module."""
         if self.lstm is not None:
             self.lstm.reset_parameters()
@@ -94,3 +98,58 @@ class JumpingKnowledge(torch.nn.Module):
             return (f'{self.__class__.__name__}({self.mode}, '
                     f'channels={self.channels}, layers={self.num_layers})')
         return f'{self.__class__.__name__}({self.mode})'
+
+
+class HeteroJumpingKnowledge(torch.nn.Module):
+    r"""A heterogeneous version of the :class:`JumpingKnowledge` module.
+
+    Args:
+        types (List[str]): The keys of the input dictionary.
+        mode (str): The aggregation scheme to use
+            (:obj:`"cat"`, :obj:`"max"` or :obj:`"lstm"`).
+        channels (int, optional): The number of channels per representation.
+            Needs to be only set for LSTM-style aggregation.
+            (default: :obj:`None`)
+        num_layers (int, optional): The number of layers to aggregate. Needs to
+            be only set for LSTM-style aggregation. (default: :obj:`None`)
+    """
+    def __init__(
+        self,
+        types: List[str],
+        mode: str,
+        channels: Optional[int] = None,
+        num_layers: Optional[int] = None,
+    ) -> None:
+        super().__init__()
+
+        self.mode = mode.lower()
+
+        self.jk_dict = torch.nn.ModuleDict({
+            key:
+            JumpingKnowledge(mode, channels, num_layers)
+            for key in types
+        })
+
+    def reset_parameters(self) -> None:
+        r"""Resets all learnable parameters of the module."""
+        for jk in self.jk_dict.values():
+            jk.reset_parameters()
+
+    def forward(self, xs_dict: Dict[str, List[Tensor]]) -> Dict[str, Tensor]:
+        r"""Forward pass.
+
+        Args:
+            xs_dict (Dict[str, List[torch.Tensor]]): A dictionary holding a
+                list of layer-wise representation for each type.
+        """
+        return {key: jk(xs_dict[key]) for key, jk in self.jk_dict.items()}
+
+    def __repr__(self):
+        if self.mode == 'lstm':
+            jk = next(iter(self.jk_dict.values()))
+            return (f'{self.__class__.__name__}('
+                    f'num_types={len(self.jk_dict)}, '
+                    f'mode={self.mode}, channels={jk.channels}, '
+                    f'layers={jk.num_layers})')
+        return (f'{self.__class__.__name__}(num_types={len(self.jk_dict)}, '
+                f'mode={self.mode})')


Test patch:
diff --git a/test/nn/models/test_jumping_knowledge.py b/test/nn/models/test_jumping_knowledge.py
index 0a03d8e42..1ffe8b2fc 100644
--- a/test/nn/models/test_jumping_knowledge.py
+++ b/test/nn/models/test_jumping_knowledge.py
@@ -1,6 +1,6 @@
 import torch
 
-from torch_geometric.nn import JumpingKnowledge
+from torch_geometric.nn import HeteroJumpingKnowledge, JumpingKnowledge
 from torch_geometric.testing import is_full_test
 
 
@@ -38,3 +38,55 @@ def test_jumping_knowledge():
     if is_full_test():
         jit = torch.jit.script(model)
         assert torch.allclose(jit(xs), out)
+
+
+def test_hetero_jumping_knowledge():
+    num_nodes, channels, num_layers = 100, 17, 5
+
+    types = ["author", "paper"]
+    xs_dict = {
+        key: [torch.randn(num_nodes, channels) for _ in range(num_layers)]
+        for key in types
+    }
+
+    model = HeteroJumpingKnowledge(types, mode='cat')
+    model.reset_parameters()
+    assert str(model) == 'HeteroJumpingKnowledge(num_types=2, mode=cat)'
+
+    out_dict = model(xs_dict)
+    for out in out_dict.values():
+        assert out.size() == (num_nodes, channels * num_layers)
+
+    if is_full_test():
+        jit = torch.jit.script(model)
+        jit_out = jit(xs_dict)
+        for key in types:
+            assert torch.allclose(jit_out[key], out_dict[key])
+
+    model = HeteroJumpingKnowledge(types, mode='max')
+    assert str(model) == 'HeteroJumpingKnowledge(num_types=2, mode=max)'
+
+    out_dict = model(xs_dict)
+    for out in out_dict.values():
+        assert out.size() == (num_nodes, channels)
+
+    if is_full_test():
+        jit = torch.jit.script(model)
+        jit_out = jit(xs_dict)
+        for key in types:
+            assert torch.allclose(jit_out[key], out_dict[key])
+
+    model = HeteroJumpingKnowledge(types, mode='lstm', channels=channels,
+                                   num_layers=num_layers)
+    assert str(model) == (f'HeteroJumpingKnowledge(num_types=2, mode=lstm, '
+                          f'channels={channels}, layers={num_layers})')
+
+    out_dict = model(xs_dict)
+    for out in out_dict.values():
+        assert out.size() == (num_nodes, channels)
+
+    if is_full_test():
+        jit = torch.jit.script(model)
+        jit_out = jit(xs_dict)
+        for key in types:
+            assert torch.allclose(jit_out[key], out_dict[key])