pyg-nightly 2.7.0.dev20250406__py3-none-any.whl → 2.7.0.dev20250408__py3-none-any.whl

{pyg_nightly-2.7.0.dev20250406.dist-info → pyg_nightly-2.7.0.dev20250408.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyg-nightly
-Version: 2.7.0.dev20250406
+Version: 2.7.0.dev20250408
 Summary: Graph Neural Network Library for PyTorch
 Keywords: deep-learning,pytorch,geometric-deep-learning,graph-neural-networks,graph-convolutional-networks
 Author-email: Matthias Fey <matthias@pyg.org>
@@ -413,18 +413,6 @@ These approaches have been implemented in PyG, and can benefit from the above GN
 
 PyG is available for Python 3.9 to Python 3.12.
 
-### Anaconda
-
-You can now install PyG via [Anaconda](https://anaconda.org/pyg/pyg) for all major OS/PyTorch/CUDA combinations 🤗
-If you have not yet installed PyTorch, install it via `conda` as described in the [official PyTorch documentation](https://pytorch.org/get-started/locally/).
-Given that you have PyTorch installed (`>=1.8.0`), simply run
-
-```
-conda install pyg -c pyg
-```
-
-### PyPi
-
 From **PyG 2.3** onwards, you can install and use PyG **without any external library** required except for PyTorch.
 For this, simply run
 
@@ -448,28 +436,28 @@ We recommend to start with a minimal installation, and install additional depend
 
 For ease of installation of these extensions, we provide `pip` wheels for all major OS/PyTorch/CUDA combinations, see [here](https://data.pyg.org/whl).
 
-#### PyTorch 2.5
+#### PyTorch 2.6
 
-To install the binaries for PyTorch 2.5.0, simply run
+To install the binaries for PyTorch 2.6.0, simply run
 
 ```
-pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.5.0+${CUDA}.html
+pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.6.0+${CUDA}.html
 ```
 
-where `${CUDA}` should be replaced by either `cpu`, `cu118`, `cu121`, or `cu124` depending on your PyTorch installation.
+where `${CUDA}` should be replaced by either `cpu`, `cu118`, `cu124`, or `cu126` depending on your PyTorch installation.
 
-|             | `cpu` | `cu118` | `cu121` | `cu124` |
+|             | `cpu` | `cu118` | `cu124` | `cu126` |
 | ----------- | ----- | ------- | ------- | ------- |
 | **Linux**   | ✅    | ✅      | ✅      | ✅      |
 | **Windows** | ✅    | ✅      | ✅      | ✅      |
 | **macOS**   | ✅    |         |         |         |
 
-#### PyTorch 2.4
+#### PyTorch 2.5
 
-To install the binaries for PyTorch 2.4.0, simply run
+To install the binaries for PyTorch 2.5.0/2.5.1, simply run
 
 ```
-pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.4.0+${CUDA}.html
+pip install pyg_lib torch_scatter torch_sparse torch_cluster torch_spline_conv -f https://data.pyg.org/whl/torch-2.5.0+${CUDA}.html
 ```
 
 where `${CUDA}` should be replaced by either `cpu`, `cu118`, `cu121`, or `cu124` depending on your PyTorch installation.
@@ -480,7 +468,7 @@ where `${CUDA}` should be replaced by either `cpu`, `cu118`, `cu121`, or `cu124`
 | **Windows** | ✅    | ✅      | ✅      | ✅      |
 | **macOS**   | ✅    |         |         |         |
 
-**Note:** Binaries of older versions are also provided for PyTorch 1.4.0, PyTorch 1.5.0, PyTorch 1.6.0, PyTorch 1.7.0/1.7.1, PyTorch 1.8.0/1.8.1, PyTorch 1.9.0, PyTorch 1.10.0/1.10.1/1.10.2, PyTorch 1.11.0, PyTorch 1.12.0/1.12.1, PyTorch 1.13.0/1.13.1, PyTorch 2.0.0/2.0.1, PyTorch 2.1.0/2.1.1/2.1.2, PyTorch 2.2.0/2.2.1/2.2.2, and PyTorch 2.3.0/2.3.1 (following the same procedure).
+**Note:** Binaries of older versions are also provided for PyTorch 1.4.0, PyTorch 1.5.0, PyTorch 1.6.0, PyTorch 1.7.0/1.7.1, PyTorch 1.8.0/1.8.1, PyTorch 1.9.0, PyTorch 1.10.0/1.10.1/1.10.2, PyTorch 1.11.0, PyTorch 1.12.0/1.12.1, PyTorch 1.13.0/1.13.1, PyTorch 2.0.0/2.0.1, PyTorch 2.1.0/2.1.1/2.1.2, PyTorch 2.2.0/2.2.1/2.2.2, PyTorch 2.3.0/2.3.1, and PyTorch 2.4.0/2.4.1 (following the same procedure).
 **For older versions, you might need to explicitly specify the latest supported version number** or install via `pip install --no-index` in order to prevent a manual installation from source.
 You can look up the latest supported version number [here](https://data.pyg.org/whl).
 

{pyg_nightly-2.7.0.dev20250406.dist-info → pyg_nightly-2.7.0.dev20250408.dist-info}/RECORD

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=EFUlgJy_cHoHOgqO8KCynWIfRJFW8DFqG7O5v9DFOzI,1978
+torch_geometric/__init__.py,sha256=rucRaUNeMS0eQLE0b6LFcLimKD5JwNS4_nYsmtC3BMU,1978
 torch_geometric/_compile.py,sha256=f-WQeH4VLi5Hn9lrgztFUCSrN_FImjhQa6BxFzcYC38,1338
 torch_geometric/_onnx.py,sha256=V9ffrIKSqhDw6xUZ12lkuSfNs48cQp2EeJ6Z19GfnVw,349
 torch_geometric/backend.py,sha256=lVaf7aLoVaB3M-UcByUJ1G4T4FOK6LXAg0CF4W3E8jo,1575
@@ -188,7 +188,7 @@ torch_geometric/distributed/event_loop.py,sha256=wr3iwMYEWOGkBlvC5huD2k5YxisaGE9
 torch_geometric/distributed/local_feature_store.py,sha256=CLW37RN0ouDufEs2tY9d2nLLvpxubiT6zgW3LIHAU8k,19058
 torch_geometric/distributed/local_graph_store.py,sha256=wNHXSS824Kk2HynbtWFXx-W4pl97UUBv6qFHAVqO3W4,8445
 torch_geometric/distributed/partition.py,sha256=X0BleuY0ROlUtVXKvvz814pJwglZQ2_6OiMi1K0Hhvo,14731
-torch_geometric/distributed/rpc.py,sha256=t0Ts4tzUE0LQyBr71i2iBjQDLe9NWkmVRf7C4xOllJc,5753
+torch_geometric/distributed/rpc.py,sha256=rJqiVR6Vbb2mpyVSC0Y5tPApqP-b1ck1Uq3IQpCsNSw,5737
 torch_geometric/distributed/utils.py,sha256=FGrr3qw7hx7EQaIjjqasurloCFJ9q_0jt8jdSIUjBeM,6567
 torch_geometric/explain/__init__.py,sha256=pRxVB33zsxhED1StRWdHboQWh3e06__g9N298Hzi42Y,359
 torch_geometric/explain/config.py,sha256=_0j67NAwPwjrWHPncNywCT-oKyMiryJNxufxVN1BFlM,7834
@@ -200,7 +200,7 @@ torch_geometric/explain/algorithm/base.py,sha256=wwJcREUFKDLFUDjRa9o4X3DWqQgMvhS
 torch_geometric/explain/algorithm/captum.py,sha256=k6hNgC5Kn9lVirOYVJzej8-hRuf5C2mPFUXFLd2wWsY,12857
 torch_geometric/explain/algorithm/captum_explainer.py,sha256=oz-c40hvdzii4_chEQPHzQo_dFjHr9HLuJhDLsqRIVU,7346
 torch_geometric/explain/algorithm/dummy_explainer.py,sha256=jvcVQmfngmUWgoKa5p7CXzju2HM5D5DfieJhZW3gbLc,2872
-torch_geometric/explain/algorithm/gnn_explainer.py,sha256=TRGwaKYn9nLn3fp0rSSzeGs9uHj2rZzfomMseDfq8O8,12454
+torch_geometric/explain/algorithm/gnn_explainer.py,sha256=iu45fGWdd4c6wNczWEAT-29HCAz7ncuoaS6cpx-xDJM,24660
 torch_geometric/explain/algorithm/graphmask_explainer.py,sha256=T2B081dK-JSpaQmutnkQd5xF3JF49_dPZCOgwqIKJDk,21367
 torch_geometric/explain/algorithm/pg_explainer.py,sha256=zPsl0tT9ISSWK1xP1KKpe1ZjUarhSB736WTtqwcmDIo,10372
 torch_geometric/explain/algorithm/utils.py,sha256=eh0ARPG41V7piVw5jdMYpV0p7WjTlpehnY-bWqPV_zg,2564
@@ -290,7 +290,7 @@ torch_geometric/loader/temporal_dataloader.py,sha256=AQ2QFeiXKbPp6I8sUeE8H7br-1_
 torch_geometric/loader/utils.py,sha256=f27mczQ7fEP2HpTsJGJxKS0slPu0j8zTba3jP8ViNck,14901
 torch_geometric/loader/zip_loader.py,sha256=3lt10fD15Rxm1WhWzypswGzCEwUz4h8OLCD1nE15yNg,3843
 torch_geometric/metrics/__init__.py,sha256=3krvDobW6vV5yHTjq2S2pmOXxNfysNG26muq7z48e94,699
-torch_geometric/metrics/link_pred.py,sha256=t2YHbEYc8Jbj_4Sb-Wdk5T5uzsSErpjBpUiSqOSf-NM,30729
+torch_geometric/metrics/link_pred.py,sha256=dtaI39JB-WqE1B-raiElns6xySRwmkbb9izbcyt6xHI,30886
 torch_geometric/nn/__init__.py,sha256=kQHHHUxFDht2ztD-XFQuv98TvC8MdodaFsIjAvltJBw,874
 torch_geometric/nn/data_parallel.py,sha256=lDAxRi83UNuzAQSj3eu9K2sQheOIU6wqR5elS6oDs90,4764
 torch_geometric/nn/encoding.py,sha256=QNjwWczYExZ1wRGBmpuqYbn6tB7NC4BU-DEgzjhcZqw,3115
@@ -441,7 +441,7 @@ torch_geometric/nn/models/g_retriever.py,sha256=CdSOasnPiMvq5AjduNTpz-LIZiNp3X0x
 torch_geometric/nn/models/git_mol.py,sha256=Wc6Hx6RDDR7sDWRWHfA5eK9e9gFsrTZ9OLmpMfoj3pE,12676
 torch_geometric/nn/models/glem.py,sha256=sT0XM4klVlci9wduvUoXupATUw9p25uXtaJBrmv3yvs,16431
 torch_geometric/nn/models/gnnff.py,sha256=15dkiLgy0LmH1hnUrpeoHioIp4BPTfjpVATpnGRt9E0,7860
-torch_geometric/nn/models/gpse.py,sha256=my-KIw_Ov8o0pXSCyh43NZRBAW95TFfmBgxzSimx8-A,42680
+torch_geometric/nn/models/gpse.py,sha256=Fwldw9N3axV--BcSnI4im1sy1r87a5mAXZAXHu_6k2Y,41932
 torch_geometric/nn/models/graph_mixer.py,sha256=mthMeCOikR8gseEsu4oJ3Cd9C35zHSv1p32ROwnG-6s,9246
 torch_geometric/nn/models/graph_unet.py,sha256=N8TSmJo8AlbZjjcame0xW_jZvMOirL5ahw6qv5Yjpbs,5586
 torch_geometric/nn/models/jumping_knowledge.py,sha256=9JR2EoViXKjcDSLb8tjJm-UHfv1mQCJvZAAEsYa0Ocw,5496
@@ -571,7 +571,7 @@ torch_geometric/transforms/remove_duplicated_edges.py,sha256=EMX6E1R_gXFoXwBqMqt
 torch_geometric/transforms/remove_isolated_nodes.py,sha256=Q89b73es1tPsAmTdS7tWTIM7JcPUpL37v3EZTAd25Fc,2449
 torch_geometric/transforms/remove_self_loops.py,sha256=JfoooSnTO2KPXXuC3KWGhLS0tqr4yiXOl3A0sVv2riM,1221
 torch_geometric/transforms/remove_training_classes.py,sha256=GMCZwI_LYo2ZF29DABZXeuM0Sn2i3twx_V3KBUGu2As,932
-torch_geometric/transforms/rooted_subgraph.py,sha256=exQ-6bRePiH44o7f_VoBnkyj79PfGjH_hyAoolTIih8,6509
+torch_geometric/transforms/rooted_subgraph.py,sha256=2P_2NSRcDNzfeR0zYScCaqZDkAvC0t3Ivr_lINO9abc,6504
 torch_geometric/transforms/sample_points.py,sha256=UfD44528J7SKH0I2_4ELM1A4hKKHCIDeMV6UzBbQAVU,2280
 torch_geometric/transforms/sign.py,sha256=bZUvUm9fMGXcYkI1GwPOW5ZC1QFu84vPBKpFnYxz2nA,2329
 torch_geometric/transforms/spherical.py,sha256=nU7h4IFw69JqUwRqaweVEBegHZWPOHDoTYYVRMzIZ7U,2320
@@ -592,7 +592,7 @@ torch_geometric/utils/_grid.py,sha256=1coutST2TMV9TSQcmpXze0GIK9odzZ9wBtbKs6u26D
 torch_geometric/utils/_homophily.py,sha256=1nXxGUATFPB3icEGpvEWUiuYbjU9gDGtlWpuLbtWhJk,5090
 torch_geometric/utils/_index_sort.py,sha256=FTJacmOsqgsyof7MJFHlVVdXhHOjR0j7siTb0UZ-YT0,1283
 torch_geometric/utils/_lexsort.py,sha256=chMEJJRXqfE6-K4vrVszdr3c338EhMZyi0Q9IEJD3p0,1403
-torch_geometric/utils/_negative_sampling.py,sha256=G4O572zAQgQQlVMz6ihhE13HFKEekLLYVXcYp4ZSdcQ,15521
+torch_geometric/utils/_negative_sampling.py,sha256=jxsmpryeoTT8qQrvIH11MgyhgoWzvqPGRAcVyU85VCU,15494
 torch_geometric/utils/_normalize_edge_index.py,sha256=H6DY-Dzi1Psr3igG_nb0U3ZPNZz-BBDntO2iuA8FtzA,1682
 torch_geometric/utils/_normalized_cut.py,sha256=uwVJkl-Q0tpY-w0nvcHajcQYcqFh1oDOf55XELdjJBU,1167
 torch_geometric/utils/_one_hot.py,sha256=vXC7l7zudYRZIwWv6mT-Biuk2zKELyqteJXLynPocPM,1404
@@ -636,7 +636,7 @@ torch_geometric/utils/undirected.py,sha256=H_nfpI0_WluOG6VfjPyldvcjL4w5USAKWu2x5
 torch_geometric/visualization/__init__.py,sha256=PyR_4K5SafsJrBr6qWrkjKr6GBL1b7FtZybyXCDEVwY,154
 torch_geometric/visualization/graph.py,sha256=ZuLPL92yGRi7lxlqsUPwL_EVVXF7P2kMcveTtW79vpA,4784
 torch_geometric/visualization/influence.py,sha256=CWMvuNA_Nf1sfbJmQgn58yS4OFpeKXeZPe7kEuvkUBw,477
-pyg_nightly-2.7.0.dev20250406.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
-pyg_nightly-2.7.0.dev20250406.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-pyg_nightly-2.7.0.dev20250406.dist-info/METADATA,sha256=csAfUo5zCWohsFtqQsynpzuSflZRg7_4f9DB3JWSVWE,63021
-pyg_nightly-2.7.0.dev20250406.dist-info/RECORD,,
+pyg_nightly-2.7.0.dev20250408.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
+pyg_nightly-2.7.0.dev20250408.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+pyg_nightly-2.7.0.dev20250408.dist-info/METADATA,sha256=OX81Kse1GucsZCReO0E00CJl4rrtT9HhHEkNEbJU6KY,62652
+pyg_nightly-2.7.0.dev20250408.dist-info/RECORD,,

torch_geometric/__init__.py

@@ -31,7 +31,7 @@ from .lazy_loader import LazyLoader
 contrib = LazyLoader('contrib', globals(), 'torch_geometric.contrib')
 graphgym = LazyLoader('graphgym', globals(), 'torch_geometric.graphgym')
 
-__version__ = '2.7.0.dev20250406'
+__version__ = '2.7.0.dev20250408'
 
 __all__ = [
     'Index',

torch_geometric/distributed/rpc.py

@@ -144,7 +144,7 @@ _rpc_call_pool: Dict[int, RPCCallBase] = {}
 @rpc_require_initialized
 def rpc_register(call: RPCCallBase) -> int:
     r"""Registers a call for RPC requests."""
-    global _rpc_call_id, _rpc_call_pool
+    global _rpc_call_id
 
     with _rpc_call_lock:
         call_id = _rpc_call_id

torch_geometric/explain/algorithm/gnn_explainer.py

@@ -1,14 +1,24 @@
 from math import sqrt
-from typing import Optional, Tuple, Union
+from typing import Dict, Optional, Tuple, Union, overload
 
 import torch
 from torch import Tensor
 from torch.nn.parameter import Parameter
 
-from torch_geometric.explain import ExplainerConfig, Explanation, ModelConfig
+from torch_geometric.explain import (
+    ExplainerConfig,
+    Explanation,
+    HeteroExplanation,
+    ModelConfig,
+)
 from torch_geometric.explain.algorithm import ExplainerAlgorithm
-from torch_geometric.explain.algorithm.utils import clear_masks, set_masks
+from torch_geometric.explain.algorithm.utils import (
+    clear_masks,
+    set_hetero_masks,
+    set_masks,
+)
 from torch_geometric.explain.config import MaskType, ModelMode, ModelTaskLevel
+from torch_geometric.typing import EdgeType, NodeType
 
 
 class GNNExplainer(ExplainerAlgorithm):
@@ -69,7 +79,9 @@ class GNNExplainer(ExplainerAlgorithm):
 
         self.node_mask = self.hard_node_mask = None
         self.edge_mask = self.hard_edge_mask = None
+        self.is_hetero = False
 
+    @overload
     def forward(
         self,
         model: torch.nn.Module,
@@ -80,30 +92,87 @@ class GNNExplainer(ExplainerAlgorithm):
         index: Optional[Union[int, Tensor]] = None,
         **kwargs,
     ) -> Explanation:
-        if isinstance(x, dict) or isinstance(edge_index, dict):
-            raise ValueError(f"Heterogeneous graphs not yet supported in "
-                             f"'{self.__class__.__name__}'")
+        ...
 
-        self._train(model, x, edge_index, target=target, index=index, **kwargs)
-
-        node_mask = self._post_process_mask(
-            self.node_mask,
-            self.hard_node_mask,
-            apply_sigmoid=True,
-        )
-        edge_mask = self._post_process_mask(
-            self.edge_mask,
-            self.hard_edge_mask,
-            apply_sigmoid=True,
-        )
+    @overload
+    def forward(
+        self,
+        model: torch.nn.Module,
+        x: Dict[NodeType, Tensor],
+        edge_index: Dict[EdgeType, Tensor],
+        *,
+        target: Tensor,
+        index: Optional[Union[int, Tensor]] = None,
+        **kwargs,
+    ) -> HeteroExplanation:
+        ...
 
+    def forward(
+        self,
+        model: torch.nn.Module,
+        x: Union[Tensor, Dict[NodeType, Tensor]],
+        edge_index: Union[Tensor, Dict[EdgeType, Tensor]],
+        *,
+        target: Tensor,
+        index: Optional[Union[int, Tensor]] = None,
+        **kwargs,
+    ) -> Union[Explanation, HeteroExplanation]:
+        self.is_hetero = isinstance(x, dict)
+        self._train(model, x, edge_index, target=target, index=index, **kwargs)
+        explanation = self._create_explanation()
         self._clean_model(model)
+        return explanation
+
+    def _create_explanation(self) -> Union[Explanation, HeteroExplanation]:
+        """Create an explanation object from the current masks."""
+        if self.is_hetero:
+            # For heterogeneous graphs, process each type separately
+            node_mask_dict = {}
+            edge_mask_dict = {}
+
+            for node_type, mask in self.node_mask.items():
+                if mask is not None:
+                    node_mask_dict[node_type] = self._post_process_mask(
+                        mask,
+                        self.hard_node_mask[node_type],
+                        apply_sigmoid=True,
+                    )
+
+            for edge_type, mask in self.edge_mask.items():
+                if mask is not None:
+                    edge_mask_dict[edge_type] = self._post_process_mask(
+                        mask,
+                        self.hard_edge_mask[edge_type],
+                        apply_sigmoid=True,
+                    )
+
+            # Create heterogeneous explanation
+            explanation = HeteroExplanation()
+            explanation.set_value_dict('node_mask', node_mask_dict)
+            explanation.set_value_dict('edge_mask', edge_mask_dict)
 
-        return Explanation(node_mask=node_mask, edge_mask=edge_mask)
+        else:
+            # For homogeneous graphs, process single masks
+            node_mask = self._post_process_mask(
+                self.node_mask,
+                self.hard_node_mask,
+                apply_sigmoid=True,
+            )
+            edge_mask = self._post_process_mask(
+                self.edge_mask,
+                self.hard_edge_mask,
+                apply_sigmoid=True,
+            )
+
+            # Create homogeneous explanation
+            explanation = Explanation(node_mask=node_mask, edge_mask=edge_mask)
+
+        return explanation
 
     def supports(self) -> bool:
         return True
 
+    @overload
     def _train(
         self,
         model: torch.nn.Module,
@@ -113,57 +182,222 @@ class GNNExplainer(ExplainerAlgorithm):
         target: Tensor,
         index: Optional[Union[int, Tensor]] = None,
         **kwargs,
-    ):
+    ) -> None:
+        ...
+
+    @overload
+    def _train(
+        self,
+        model: torch.nn.Module,
+        x: Dict[NodeType, Tensor],
+        edge_index: Dict[EdgeType, Tensor],
+        *,
+        target: Tensor,
+        index: Optional[Union[int, Tensor]] = None,
+        **kwargs,
+    ) -> None:
+        ...
+
+    def _train(
+        self,
+        model: torch.nn.Module,
+        x: Union[Tensor, Dict[NodeType, Tensor]],
+        edge_index: Union[Tensor, Dict[EdgeType, Tensor]],
+        *,
+        target: Tensor,
+        index: Optional[Union[int, Tensor]] = None,
+        **kwargs,
+    ) -> None:
+        # Initialize masks based on input type
         self._initialize_masks(x, edge_index)
 
-        parameters = []
-        if self.node_mask is not None:
-            parameters.append(self.node_mask)
-        if self.edge_mask is not None:
-            set_masks(model, self.edge_mask, edge_index, apply_sigmoid=True)
-            parameters.append(self.edge_mask)
+        # Collect parameters for optimization
+        parameters = self._collect_parameters(model, edge_index)
 
+        # Create optimizer
         optimizer = torch.optim.Adam(parameters, lr=self.lr)
 
+        # Training loop
         for i in range(self.epochs):
             optimizer.zero_grad()
 
-            h = x if self.node_mask is None else x * self.node_mask.sigmoid()
-            y_hat, y = model(h, edge_index, **kwargs), target
+            # Forward pass with masked inputs
+            y_hat = self._forward_with_masks(model, x, edge_index, **kwargs)
+            y = target
 
+            # Handle index if provided
             if index is not None:
                 y_hat, y = y_hat[index], y[index]
 
+            # Calculate loss
             loss = self._loss(y_hat, y)
 
+            # Backward pass
             loss.backward()
             optimizer.step()
 
-            # In the first iteration, we collect the nodes and edges that are
-            # involved into making the prediction. These are all the nodes and
-            # edges with gradient != 0 (without regularization applied).
-            if i == 0 and self.node_mask is not None:
-                if self.node_mask.grad is None:
-                    raise ValueError("Could not compute gradients for node "
-                                     "features. Please make sure that node "
-                                     "features are used inside the model or "
-                                     "disable it via `node_mask_type=None`.")
-                self.hard_node_mask = self.node_mask.grad != 0.0
-            if i == 0 and self.edge_mask is not None:
-                if self.edge_mask.grad is None:
-                    raise ValueError("Could not compute gradients for edges. "
-                                     "Please make sure that edges are used "
-                                     "via message passing inside the model or "
-                                     "disable it via `edge_mask_type=None`.")
-                self.hard_edge_mask = self.edge_mask.grad != 0.0
-
-    def _initialize_masks(self, x: Tensor, edge_index: Tensor):
+            # In the first iteration, collect gradients to identify important
+            # nodes/edges
+            if i == 0:
+                self._collect_gradients()
+
+    def _collect_parameters(self, model, edge_index):
+        """Collect parameters for optimization."""
+        parameters = []
+
+        if self.is_hetero:
+            # For heterogeneous graphs, collect parameters from all types
+            for mask in self.node_mask.values():
+                if mask is not None:
+                    parameters.append(mask)
+            if any(v is not None for v in self.edge_mask.values()):
+                set_hetero_masks(model, self.edge_mask, edge_index)
+            for mask in self.edge_mask.values():
+                if mask is not None:
+                    parameters.append(mask)
+        else:
+            # For homogeneous graphs, collect single parameters
+            if self.node_mask is not None:
+                parameters.append(self.node_mask)
+            if self.edge_mask is not None:
+                set_masks(model, self.edge_mask, edge_index,
+                          apply_sigmoid=True)
+                parameters.append(self.edge_mask)
+
+        return parameters
+
+    @overload
+    def _forward_with_masks(
+        self,
+        model: torch.nn.Module,
+        x: Tensor,
+        edge_index: Tensor,
+        **kwargs,
+    ) -> Tensor:
+        ...
+
+    @overload
+    def _forward_with_masks(
+        self,
+        model: torch.nn.Module,
+        x: Dict[NodeType, Tensor],
+        edge_index: Dict[EdgeType, Tensor],
+        **kwargs,
+    ) -> Tensor:
+        ...
+
+    def _forward_with_masks(
+        self,
+        model: torch.nn.Module,
+        x: Union[Tensor, Dict[NodeType, Tensor]],
+        edge_index: Union[Tensor, Dict[EdgeType, Tensor]],
+        **kwargs,
+    ) -> Tensor:
+        """Forward pass with masked inputs."""
+        if self.is_hetero:
+            # Apply masks to heterogeneous inputs
+            h_dict = {}
+            for node_type, features in x.items():
+                if node_type in self.node_mask and self.node_mask[
+                        node_type] is not None:
+                    h_dict[node_type] = features * self.node_mask[
+                        node_type].sigmoid()
+                else:
+                    h_dict[node_type] = features
+
+            # Forward pass with masked features
+            return model(h_dict, edge_index, **kwargs)
+        else:
+            # Apply mask to homogeneous input
+            h = x if self.node_mask is None else x * self.node_mask.sigmoid()
+
+            # Forward pass with masked features
+            return model(h, edge_index, **kwargs)
+
+    def _initialize_masks(
+        self,
+        x: Union[Tensor, Dict[NodeType, Tensor]],
+        edge_index: Union[Tensor, Dict[EdgeType, Tensor]],
+    ) -> None:
         node_mask_type = self.explainer_config.node_mask_type
         edge_mask_type = self.explainer_config.edge_mask_type
 
-        device = x.device
-        (N, F), E = x.size(), edge_index.size(1)
+        if self.is_hetero:
+            # Initialize dictionaries for heterogeneous masks
+            self.node_mask = {}
+            self.hard_node_mask = {}
+            self.edge_mask = {}
+            self.hard_edge_mask = {}
+
+            # Initialize node masks for each node type
+            for node_type, features in x.items():
+                device = features.device
+                N, F = features.size()
+                self._initialize_node_mask(node_mask_type, node_type, N, F,
+                                           device)
+
+            # Initialize edge masks for each edge type
+            for edge_type, indices in edge_index.items():
+                device = indices.device
+                E = indices.size(1)
+                N = max(indices.max().item() + 1,
+                        max(feat.size(0) for feat in x.values()))
+                self._initialize_edge_mask(edge_mask_type, edge_type, E, N,
+                                           device)
+        else:
+            # Initialize masks for homogeneous graph
+            device = x.device
+            (N, F), E = x.size(), edge_index.size(1)
+
+            # Initialize homogeneous node and edge masks
+            self._initialize_homogeneous_masks(node_mask_type, edge_mask_type,
+                                               N, F, E, device)
+
+    def _initialize_node_mask(
+        self,
+        node_mask_type,
+        node_type,
+        N,
+        F,
+        device,
+    ) -> None:
+        """Initialize node mask for a specific node type."""
+        std = 0.1
+        if node_mask_type is None:
+            self.node_mask[node_type] = None
+            self.hard_node_mask[node_type] = None
+        elif node_mask_type == MaskType.object:
+            self.node_mask[node_type] = Parameter(
+                torch.randn(N, 1, device=device) * std)
+            self.hard_node_mask[node_type] = None
+        elif node_mask_type == MaskType.attributes:
+            self.node_mask[node_type] = Parameter(
+                torch.randn(N, F, device=device) * std)
+            self.hard_node_mask[node_type] = None
+        elif node_mask_type == MaskType.common_attributes:
+            self.node_mask[node_type] = Parameter(
+                torch.randn(1, F, device=device) * std)
+            self.hard_node_mask[node_type] = None
+        else:
+            raise ValueError(f"Invalid node mask type: {node_mask_type}")
+
+    def _initialize_edge_mask(self, edge_mask_type, edge_type, E, N, device):
+        """Initialize edge mask for a specific edge type."""
+        if edge_mask_type is None:
+            self.edge_mask[edge_type] = None
+            self.hard_edge_mask[edge_type] = None
+        elif edge_mask_type == MaskType.object:
+            std = torch.nn.init.calculate_gain('relu') * sqrt(2.0 / (2 * N))
+            self.edge_mask[edge_type] = Parameter(
+                torch.randn(E, device=device) * std)
+            self.hard_edge_mask[edge_type] = None
+        else:
+            raise ValueError(f"Invalid edge mask type: {edge_mask_type}")
 
+    def _initialize_homogeneous_masks(self, node_mask_type, edge_mask_type, N,
+                                      F, E, device):
+        """Initialize masks for homogeneous graph."""
+        # Initialize node mask
         std = 0.1
         if node_mask_type is None:
             self.node_mask = None
@@ -174,43 +408,145 @@ class GNNExplainer(ExplainerAlgorithm):
         elif node_mask_type == MaskType.common_attributes:
             self.node_mask = Parameter(torch.randn(1, F, device=device) * std)
         else:
-            assert False
+            raise ValueError(f"Invalid node mask type: {node_mask_type}")
 
+        # Initialize edge mask
         if edge_mask_type is None:
             self.edge_mask = None
         elif edge_mask_type == MaskType.object:
             std = torch.nn.init.calculate_gain('relu') * sqrt(2.0 / (2 * N))
             self.edge_mask = Parameter(torch.randn(E, device=device) * std)
         else:
-            assert False
+            raise ValueError(f"Invalid edge mask type: {edge_mask_type}")
+
+    def _collect_gradients(self) -> None:
+        if self.is_hetero:
+            self._collect_hetero_gradients()
+        else:
+            self._collect_homo_gradients()
+
+    def _collect_hetero_gradients(self):
+        """Collect gradients for heterogeneous graph."""
+        for node_type, mask in self.node_mask.items():
+            if mask is not None:
+                if mask.grad is None:
+                    raise ValueError(
+                        f"Could not compute gradients for node masks of type "
+                        f"'{node_type}'. Please make sure that node masks are "
+                        f"used inside the model or disable it via "
+                        f"`node_mask_type=None`.")
+
+                self.hard_node_mask[node_type] = mask.grad != 0.0
+
+        for edge_type, mask in self.edge_mask.items():
+            if mask is not None:
+                if mask.grad is None:
+                    raise ValueError(
+                        f"Could not compute gradients for edge masks of type "
+                        f"'{edge_type}'. Please make sure that edge masks are "
+                        f"used inside the model or disable it via "
+                        f"`edge_mask_type=None`.")
+                self.hard_edge_mask[edge_type] = mask.grad != 0.0
+
+    def _collect_homo_gradients(self):
+        """Collect gradients for homogeneous graph."""
+        if self.node_mask is not None:
+            if self.node_mask.grad is None:
+                raise ValueError("Could not compute gradients for node "
+                                 "features. Please make sure that node "
+                                 "features are used inside the model or "
+                                 "disable it via `node_mask_type=None`.")
+            self.hard_node_mask = self.node_mask.grad != 0.0
+
+        if self.edge_mask is not None:
+            if self.edge_mask.grad is None:
+                raise ValueError("Could not compute gradients for edges. "
+                                 "Please make sure that edges are used "
+                                 "via message passing inside the model or "
+                                 "disable it via `edge_mask_type=None`.")
+            self.hard_edge_mask = self.edge_mask.grad != 0.0
 
     def _loss(self, y_hat: Tensor, y: Tensor) -> Tensor:
+        # Calculate base loss based on model configuration
+        loss = self._calculate_base_loss(y_hat, y)
+
+        # Apply regularization based on graph type
+        if self.is_hetero:
+            # Apply regularization for heterogeneous graph
+            loss = self._apply_hetero_regularization(loss)
+        else:
+            # Apply regularization for homogeneous graph
+            loss = self._apply_homo_regularization(loss)
+
+        return loss
+
+    def _calculate_base_loss(self, y_hat, y):
+        """Calculate base loss based on model configuration."""
         if self.model_config.mode == ModelMode.binary_classification:
-            loss = self._loss_binary_classification(y_hat, y)
+            return self._loss_binary_classification(y_hat, y)
         elif self.model_config.mode == ModelMode.multiclass_classification:
-            loss = self._loss_multiclass_classification(y_hat, y)
+            return self._loss_multiclass_classification(y_hat, y)
         elif self.model_config.mode == ModelMode.regression:
-            loss = self._loss_regression(y_hat, y)
+            return self._loss_regression(y_hat, y)
         else:
-            assert False
+            raise ValueError(f"Invalid model mode: {self.model_config.mode}")
+
+    def _apply_hetero_regularization(self, loss):
+        """Apply regularization for heterogeneous graph."""
+        # Apply regularization for each edge type
+        for edge_type, mask in self.edge_mask.items():
+            if (mask is not None
+                    and self.hard_edge_mask[edge_type] is not None):
+                loss = self._add_mask_regularization(
+                    loss, mask, self.hard_edge_mask[edge_type],
+                    self.coeffs['edge_size'], self.coeffs['edge_reduction'],
+                    self.coeffs['edge_ent'])
+
+        # Apply regularization for each node type
+        for node_type, mask in self.node_mask.items():
+            if (mask is not None
+                    and self.hard_node_mask[node_type] is not None):
+                loss = self._add_mask_regularization(
+                    loss, mask, self.hard_node_mask[node_type],
+                    self.coeffs['node_feat_size'],
+                    self.coeffs['node_feat_reduction'],
+                    self.coeffs['node_feat_ent'])
 
+        return loss
+
+    def _apply_homo_regularization(self, loss):
+        """Apply regularization for homogeneous graph."""
+        # Apply regularization for edge mask
         if self.hard_edge_mask is not None:
             assert self.edge_mask is not None
-            m = self.edge_mask[self.hard_edge_mask].sigmoid()
-            edge_reduce = getattr(torch, self.coeffs['edge_reduction'])
-            loss = loss + self.coeffs['edge_size'] * edge_reduce(m)
-            ent = -m * torch.log(m + self.coeffs['EPS']) - (
-                1 - m) * torch.log(1 - m + self.coeffs['EPS'])
-            loss = loss + self.coeffs['edge_ent'] * ent.mean()
+            loss = self._add_mask_regularization(loss, self.edge_mask,
+                                                 self.hard_edge_mask,
+                                                 self.coeffs['edge_size'],
+                                                 self.coeffs['edge_reduction'],
+                                                 self.coeffs['edge_ent'])
 
+        # Apply regularization for node mask
         if self.hard_node_mask is not None:
             assert self.node_mask is not None
-            m = self.node_mask[self.hard_node_mask].sigmoid()
-            node_reduce = getattr(torch, self.coeffs['node_feat_reduction'])
-            loss = loss + self.coeffs['node_feat_size'] * node_reduce(m)
-            ent = -m * torch.log(m + self.coeffs['EPS']) - (
-                1 - m) * torch.log(1 - m + self.coeffs['EPS'])
-            loss = loss + self.coeffs['node_feat_ent'] * ent.mean()
+            loss = self._add_mask_regularization(
+                loss, self.node_mask, self.hard_node_mask,
+                self.coeffs['node_feat_size'],
+                self.coeffs['node_feat_reduction'],
+                self.coeffs['node_feat_ent'])
+
+        return loss
+
+    def _add_mask_regularization(self, loss, mask, hard_mask, size_coeff,
+                                 reduction_name, ent_coeff):
+        """Add size and entropy regularization for a mask."""
+        m = mask[hard_mask].sigmoid()
+        reduce_fn = getattr(torch, reduction_name)
+        # Add size regularization
+        loss = loss + size_coeff * reduce_fn(m)
+        # Add entropy regularization
+        ent = -m * torch.log(m + self.coeffs['EPS']) - (
+            1 - m) * torch.log(1 - m + self.coeffs['EPS'])
+        loss = loss + ent_coeff * ent.mean()
 
         return loss
 

torch_geometric/metrics/link_pred.py

@@ -670,6 +670,9 @@ class LinkPredDiversity(_LinkPredMetric):
     def __init__(self, k: int, category: Tensor) -> None:
         super().__init__(k)
 
+        self.accum: Tensor
+        self.total: Tensor
+
         if WITH_TORCHMETRICS:
             self.add_state('accum', torch.tensor(0.), dist_reduce_fx='sum')
             self.add_state('total', torch.tensor(0), dist_reduce_fx='sum')
@@ -736,11 +739,14 @@ class LinkPredPersonalization(_LinkPredMetric):
         self.max_src_nodes = max_src_nodes
         self.batch_size = batch_size
 
+        self.preds: List[Tensor]
+        self.total: Tensor
+
         if WITH_TORCHMETRICS:
             self.add_state('preds', default=[], dist_reduce_fx='cat')
             self.add_state('total', torch.tensor(0), dist_reduce_fx='sum')
         else:
-            self.preds: List[Tensor] = []
+            self.preds = []
             self.register_buffer('total', torch.tensor(0), persistent=False)
 
     def update(
@@ -826,6 +832,9 @@ class LinkPredAveragePopularity(_LinkPredMetric):
     def __init__(self, k: int, popularity: Tensor) -> None:
         super().__init__(k)
 
+        self.accum: Tensor
+        self.total: Tensor
+
         if WITH_TORCHMETRICS:
             self.add_state('accum', torch.tensor(0.), dist_reduce_fx='sum')
             self.add_state('total', torch.tensor(0), dist_reduce_fx='sum')

torch_geometric/nn/models/gpse.py

@@ -379,7 +379,7 @@ class GPSE(torch.nn.Module):
 
     .. code-block:: python
 
-        from torch_geometric.nn import GPSE, GPSENodeEncoder,
+        from torch_geometric.nn import GPSE, GPSENodeEncoder
         from torch_geometric.transforms import AddGPSE
         from torch_geometric.nn.models.gpse import precompute_GPSE
 
@@ -417,13 +417,11 @@ class GPSE(torch.nn.Module):
 
         encoder = GPSENodeEncoder(dim_emb=128, dim_pe_in=32, dim_pe_out=64,
                                   expand_x=False)
-        gnn = GNN(dim_in=128, dim_out=128, num_layers=4)
+        gnn = GNN(...)
 
         for batch in loader:
-            batch = encoder(batch)
-            batch = gnn(batch)
-            # Do something with the batch, which now includes 128-dimensional
-            # node representations
+            x = encoder(batch.x, batch.pestat_GPSE)
+            out = gnn(x, batch.edge_index)
 
 
     Args:
@@ -571,8 +569,7 @@ class GPSE(torch.nn.Module):
         self.reset_parameters()
 
     def reset_parameters(self):
-        from torch_geometric.graphgym.init import init_weights
-        self.apply(init_weights)
+        pass
 
     @classmethod
     def from_pretrained(cls, name: str, root: str = 'GPSE_pretrained'):
@@ -608,6 +605,7 @@ class GPSE(torch.nn.Module):
         return model
 
     def forward(self, batch):
+        batch = batch.clone()
         for module in self.children():
             batch = module(batch)
         return batch
@@ -635,13 +633,11 @@ class GPSENodeEncoder(torch.nn.Module):
 
         encoder = GPSENodeEncoder(dim_emb=128, dim_pe_in=32, dim_pe_out=64,
                                   expand_x=False)
-        gnn = GNN(dim_in=128, dim_out=128, num_layers=4)
+        gnn = GNN(...)
 
         for batch in loader:
-            batch = encoder(batch)
-            batch = gnn(batch)
-            # Do something with the batch, which now includes 128-dimensional
-            # node representations
+            x = encoder(batch.x, batch.pestat_GPSE)
+            batch = gnn(x, batch.edge_index)
 
     Args:
         dim_emb (int): Size of final node embedding.
@@ -705,28 +701,17 @@ class GPSENodeEncoder(torch.nn.Module):
             raise ValueError(f"{self.__class__.__name__}: Does not support "
                              f"'{model_type}' encoder model.")
 
-    def forward(self, batch):
-        if not hasattr(batch, 'pestat_GPSE'):
-            raise ValueError('Precomputed "pestat_GPSE" variable is required '
-                             'for GNNNodeEncoder; either run '
-                             '`precompute_GPSE(gpse_model, dataset)` on your '
-                             'dataset or add `AddGPSE(gpse_model)` as a (pre) '
-                             'transform.')
-
-        pos_enc = batch.pestat_GPSE
-
+    def forward(self, x, pos_enc):
         pos_enc = self.dropout_be(pos_enc)
         pos_enc = self.raw_norm(pos_enc) if self.raw_norm else pos_enc
         pos_enc = self.pe_encoder(pos_enc)  # (Num nodes) x dim_pe
         pos_enc = self.dropout_ae(pos_enc)
 
         # Expand node features if needed
-        h = self.linear_x(batch.x) if self.expand_x else batch.x
+        h = self.linear_x(x) if self.expand_x else x
 
         # Concatenate final PEs to input embedding
-        batch.x = torch.cat((h, pos_enc), 1)
-
-        return batch
+        return torch.cat((h, pos_enc), 1)
 
 
 @torch.no_grad()

torch_geometric/transforms/rooted_subgraph.py

@@ -94,7 +94,7 @@ class RootedSubgraph(BaseTransform, ABC):
         arange = torch.arange(n_id.size(0), device=data.edge_index.device)
         node_map = data.edge_index.new_ones(num_nodes, num_nodes)
         node_map[n_sub_batch, n_id] = arange
-        sub_edge_index += (arange * data.num_nodes)[e_sub_batch]
+        sub_edge_index += (arange * num_nodes)[e_sub_batch]
         sub_edge_index = node_map.view(-1)[sub_edge_index]
 
         return sub_edge_index, n_id, e_id, n_sub_batch, e_sub_batch

torch_geometric/utils/_negative_sampling.py

@@ -109,10 +109,9 @@ def negative_sampling(
         idx = idx.to('cpu')
         for _ in range(3):  # Number of tries to sample negative indices.
             rnd = sample(population, sample_size, device='cpu')
-            mask = np.isin(rnd.numpy(), idx.numpy())  # type: ignore
+            mask = torch.from_numpy(np.isin(rnd.numpy(), idx.numpy())).bool()
             if neg_idx is not None:
-                mask |= np.isin(rnd, neg_idx.to('cpu'))
-            mask = torch.from_numpy(mask).to(torch.bool)
+                mask |= torch.from_numpy(np.isin(rnd, neg_idx.cpu())).bool()
             rnd = rnd[~mask].to(edge_index.device)
             neg_idx = rnd if neg_idx is None else torch.cat([neg_idx, rnd])
             if neg_idx.numel() >= num_neg_samples: