Repository: pyg-team/pytorch_geometric
Pull request #9394 (2024-06-05 00:48:58)

 [FutureWarning] Fixing the warning that appears when calling the constructor of the `LogisticRegression` class
The following warning, which appeared in three tests, has been fixed:
```
/usr/local/lib/python3.10/dist-packages/sklearn/linear_model/_logistic.py:1247: FutureWarning: 'multi_class' was 
deprecated in version 1.5 and will be removed in 1.7. From then on, it will always use 'multinomial'. Leave it to 
its default value to avoid this warning.
```


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/deep_graph_infomax.py b/torch_geometric/nn/models/deep_graph_infomax.py
index 306e905c6..e46967f96 100644
--- a/torch_geometric/nn/models/deep_graph_infomax.py
+++ b/torch_geometric/nn/models/deep_graph_infomax.py
@@ -98,7 +98,6 @@ class DeepGraphInfomax(torch.nn.Module):
         test_z: Tensor,
         test_y: Tensor,
         solver: str = 'lbfgs',
-        multi_class: str = 'auto',
         *args,
         **kwargs,
     ) -> float:
@@ -107,7 +106,7 @@ class DeepGraphInfomax(torch.nn.Module):
         """
         from sklearn.linear_model import LogisticRegression
 
-        clf = LogisticRegression(solver=solver, multi_class=multi_class, *args,
+        clf = LogisticRegression(solver=solver, *args,
                                  **kwargs).fit(train_z.detach().cpu().numpy(),
                                                train_y.detach().cpu().numpy())
         return clf.score(test_z.detach().cpu().numpy(),
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})')
diff --git a/torch_geometric/nn/models/metapath2vec.py b/torch_geometric/nn/models/metapath2vec.py
index 406d44617..f077e6f88 100644
--- a/torch_geometric/nn/models/metapath2vec.py
+++ b/torch_geometric/nn/models/metapath2vec.py
@@ -227,14 +227,13 @@ class MetaPath2Vec(torch.nn.Module):
         return pos_loss + neg_loss
 
     def test(self, train_z: Tensor, train_y: Tensor, test_z: Tensor,
-             test_y: Tensor, solver: str = "lbfgs", multi_class: str = "auto",
-             *args, **kwargs) -> float:
+             test_y: Tensor, solver: str = "lbfgs", *args, **kwargs) -> float:
         r"""Evaluates latent space quality via a logistic regression downstream
         task.
         """
         from sklearn.linear_model import LogisticRegression
 
-        clf = LogisticRegression(solver=solver, multi_class=multi_class, *args,
+        clf = LogisticRegression(solver=solver, *args,
                                  **kwargs).fit(train_z.detach().cpu().numpy(),
                                                train_y.detach().cpu().numpy())
         return clf.score(test_z.detach().cpu().numpy(),
diff --git a/torch_geometric/nn/models/node2vec.py b/torch_geometric/nn/models/node2vec.py
index 1b61bc5c7..e16c12191 100644
--- a/torch_geometric/nn/models/node2vec.py
+++ b/torch_geometric/nn/models/node2vec.py
@@ -173,7 +173,6 @@ class Node2Vec(torch.nn.Module):
         test_z: Tensor,
         test_y: Tensor,
         solver: str = 'lbfgs',
-        multi_class: str = 'auto',
         *args,
         **kwargs,
     ) -> float:
@@ -182,7 +181,7 @@ class Node2Vec(torch.nn.Module):
         """
         from sklearn.linear_model import LogisticRegression
 
-        clf = LogisticRegression(solver=solver, multi_class=multi_class, *args,
+        clf = LogisticRegression(solver=solver, *args,
                                  **kwargs).fit(train_z.detach().cpu().numpy(),
                                                train_y.detach().cpu().numpy())
         return clf.score(test_z.detach().cpu().numpy(),
diff --git a/torch_geometric/nn/to_hetero_with_bases_transformer.py b/torch_geometric/nn/to_hetero_with_bases_transformer.py
index 2bcf4d8ee..b153187cf 100644
--- a/torch_geometric/nn/to_hetero_with_bases_transformer.py
+++ b/torch_geometric/nn/to_hetero_with_bases_transformer.py
@@ -309,6 +309,24 @@ class ToHeteroWithBasesTransformer(Transformer):
 ###############################################################################
 
 
+# We make use of a post-message computation hook to inject the
+# basis re-weighting for each individual edge type.
+# This currently requires us to set `conv.fuse = False`, which leads
+# to a materialization of messages.
+def hook(module, inputs, output):
+    assert isinstance(module._edge_type, Tensor)
+    if module._edge_type.size(0) != output.size(-2):
+        raise ValueError(
+            f"Number of messages ({output.size(0)}) does not match "
+            f"with the number of original edges "
+            f"({module._edge_type.size(0)}). Does your message "
+            f"passing layer create additional self-loops? Try to "
+            f"remove them via 'add_self_loops=False'")
+    weight = module.edge_type_weight.view(-1)[module._edge_type]
+    weight = weight.view([1] * (output.dim() - 2) + [-1, 1])
+    return weight * output
+
+
 class HeteroBasisConv(torch.nn.Module):
     # A wrapper layer that applies the basis-decomposition technique to a
     # heterogeneous graph.
@@ -319,23 +337,6 @@ class HeteroBasisConv(torch.nn.Module):
         self.num_relations = num_relations
         self.num_bases = num_bases
 
-        # We make use of a post-message computation hook to inject the
-        # basis re-weighting for each individual edge type.
-        # This currently requires us to set `conv.fuse = False`, which leads
-        # to a materialization of messages.
-        def hook(module, inputs, output):
-            assert isinstance(module._edge_type, Tensor)
-            if module._edge_type.size(0) != output.size(-2):
-                raise ValueError(
-                    f"Number of messages ({output.size(0)}) does not match "
-                    f"with the number of original edges "
-                    f"({module._edge_type.size(0)}). Does your message "
-                    f"passing layer create additional self-loops? Try to "
-                    f"remove them via 'add_self_loops=False'")
-            weight = module.edge_type_weight.view(-1)[module._edge_type]
-            weight = weight.view([1] * (output.dim() - 2) + [-1, 1])
-            return weight * output
-
         params = list(module.parameters())
         device = params[0].device if len(params) > 0 else 'cpu'
 


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])
diff --git a/test/nn/test_to_hetero_with_bases_transformer.py b/test/nn/test_to_hetero_with_bases_transformer.py
index 3736ad6c6..9ecdd52b7 100644
--- a/test/nn/test_to_hetero_with_bases_transformer.py
+++ b/test/nn/test_to_hetero_with_bases_transformer.py
@@ -1,3 +1,4 @@
+import os.path as osp
 from typing import Tuple
 
 import pytest
@@ -313,4 +314,27 @@ def test_to_hetero_with_bases_on_static_graphs():
 
     assert len(out_dict) == 2
     assert out_dict['paper'].size() == (4, 100, 32)
-    assert out_dict['author'].size() == (4, 100, 32)
+
+
+def test_to_hetero_with_bases_save(tmp_path):
+    x_dict = {'paper': torch.randn(100, 16), 'author': torch.randn(100, 8)}
+    edge_index_dict = {
+        ('paper', 'cites', 'paper'):
+        torch.randint(100, (2, 200), dtype=torch.long),
+        ('paper', 'written_by', 'author'):
+        torch.randint(100, (2, 200), dtype=torch.long),
+        ('author', 'writes', 'paper'):
+        torch.randint(100, (2, 200), dtype=torch.long),
+    }
+
+    model = to_hetero_with_bases(
+        Net2(),
+        (list(x_dict.keys()), list(edge_index_dict.keys())),
+        num_bases=4,
+        in_channels={'x': 16},
+        debug=False,
+    )
+    model(x_dict, edge_index_dict)
+
+    path = osp.join(tmp_path, 'model.pt')
+    torch.save(model, path)