pyg-nightly 2.7.0.dev20250416__py3-none-any.whl → 2.7.0.dev20250418__py3-none-any.whl

{pyg_nightly-2.7.0.dev20250416.dist-info → pyg_nightly-2.7.0.dev20250418.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyg-nightly
-Version: 2.7.0.dev20250416
+Version: 2.7.0.dev20250418
 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>

{pyg_nightly-2.7.0.dev20250416.dist-info → pyg_nightly-2.7.0.dev20250418.dist-info}/RECORD

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=en3YAn8JWtZUWaMs1gx9X-Z0xZM9IfGB3fvTW53cJno,1978
+torch_geometric/__init__.py,sha256=UCc69_Z_CAoLuIzpuXALW5om7y4PXXzvqORBng1FG70,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
@@ -202,7 +202,7 @@ torch_geometric/explain/algorithm/captum_explainer.py,sha256=oz-c40hvdzii4_chEQP
 torch_geometric/explain/algorithm/dummy_explainer.py,sha256=jvcVQmfngmUWgoKa5p7CXzju2HM5D5DfieJhZW3gbLc,2872
 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/pg_explainer.py,sha256=LMlNcqSqtEP-IzYA7Xix6FoAogcrLUaEUAxDVyz2eyc,20162
 torch_geometric/explain/algorithm/utils.py,sha256=eh0ARPG41V7piVw5jdMYpV0p7WjTlpehnY-bWqPV_zg,2564
 torch_geometric/explain/metric/__init__.py,sha256=swLeuWVaM3K7UvowsH7q3BzfTq_W1vhcFY8nEP7vFPQ,301
 torch_geometric/explain/metric/basic.py,sha256=qN-cho4lxwPlw_X26svJrW5QOnw5GB3lLKf0Js_6rBE,1888
@@ -584,7 +584,7 @@ torch_geometric/transforms/to_superpixels.py,sha256=g8ysBv-ezcHn2gHucKuBtnbe-kBD
 torch_geometric/transforms/to_undirected.py,sha256=oklgrNzev7HjvVaBHwPQFo0RxcQpmcIebNbcv6vNCtY,2972
 torch_geometric/transforms/two_hop.py,sha256=XxZl3eztTjE00ZlyAIqYu36rjaRddQT-1v4AFF9VUBc,1313
 torch_geometric/transforms/virtual_node.py,sha256=FMGT6LZBH-SU2zmp76GKNqJBZ8PyS1_6Em2BbVhv8Tw,2932
-torch_geometric/utils/__init__.py,sha256=zSiljeQIG8aVXDL9Jowv6WJynfiSLt2w29XzUSu59CI,4930
+torch_geometric/utils/__init__.py,sha256=aVet2bjRvr3URikJ6LpjLATz447YuBS6FuSu5l3JwLY,4982
 torch_geometric/utils/_assortativity.py,sha256=pe2Hv5xLWhTW7dgqVWNiwDgDVMxMbliTdLeQf5Y65Ug,2347
 torch_geometric/utils/_coalesce.py,sha256=m4s_maBhib0jByQi6Cd8dazzhFVshZXLfB9aykCZT2g,6769
 torch_geometric/utils/_degree.py,sha256=FcsGx5cQdrBmoCQ4qQ2csjsTiDICP1as4x1HD9y5XVk,1017
@@ -613,7 +613,7 @@ torch_geometric/utils/augmentation.py,sha256=1F0YCuaklZ9ZbXxdFV0oOoemWvLd8p60WvF
 torch_geometric/utils/convert.py,sha256=0KEJoBOzU-w-mMQu9QYaMhUqcrGBxBmeRl0hv8NPvII,21697
 torch_geometric/utils/cross_entropy.py,sha256=ZFS5bivtzv3EV9zqgKsekmuQyoZZggPSclhl_tRNHxo,3047
 torch_geometric/utils/dropout.py,sha256=gg0rDnD4FLvBaKSoLAkZwViAQflhLefJm6_Mju5dmQs,11416
-torch_geometric/utils/embedding.py,sha256=gcWcUv46W0bZBm9puUr7GpPrdzb-PWlD9bpwqBHnA-w,1675
+torch_geometric/utils/embedding.py,sha256=b-CQ-aapEgahxSS7fuL4aNQX6GJROboV0xclZ_MmwO0,5179
 torch_geometric/utils/functions.py,sha256=orQdS_6EpzWSmBHSok3WhxCzLy9neB-cin1aTnlXY-8,703
 torch_geometric/utils/geodesic.py,sha256=-xsqE3FZU7Y9gMbucIlGJ4FM-3nk8o0AQBxIdN-QfEw,4770
 torch_geometric/utils/hetero.py,sha256=ok4uAAOyMiaeEPmvyS4DNoDwdKnLS2gmgs5WVVklxOo,5539
@@ -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.dev20250416.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
-pyg_nightly-2.7.0.dev20250416.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-pyg_nightly-2.7.0.dev20250416.dist-info/METADATA,sha256=yDodf56EgttruZas0nqEgbfnFaNHm03BHQsFi-IkPf0,62979
-pyg_nightly-2.7.0.dev20250416.dist-info/RECORD,,
+pyg_nightly-2.7.0.dev20250418.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
+pyg_nightly-2.7.0.dev20250418.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+pyg_nightly-2.7.0.dev20250418.dist-info/METADATA,sha256=j8ciN8XlpllB1ulU5n8vM4eWsUnDsDiMaGZOKoyC0Rk,62979
+pyg_nightly-2.7.0.dev20250418.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.dev20250416'
+__version__ = '2.7.0.dev20250418'
 
 __all__ = [
     'Index',

torch_geometric/explain/algorithm/pg_explainer.py

@@ -1,21 +1,26 @@
 import logging
-from typing import Optional, Union
+from typing import Dict, Optional, Tuple, Union, overload
 
 import torch
 from torch import Tensor
 from torch.nn import ReLU, Sequential
 
-from torch_geometric.explain import Explanation
+from torch_geometric.explain import Explanation, HeteroExplanation
 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 (
     ExplanationType,
     ModelMode,
     ModelTaskLevel,
 )
-from torch_geometric.nn import Linear
+from torch_geometric.nn import HANConv, HeteroConv, HGTConv, Linear
 from torch_geometric.nn.inits import reset
-from torch_geometric.utils import get_embeddings
+from torch_geometric.typing import EdgeType, NodeType
+from torch_geometric.utils import get_embeddings, get_embeddings_hetero
 
 
 class PGExplainer(ExplainerAlgorithm):
@@ -62,6 +67,13 @@ class PGExplainer(ExplainerAlgorithm):
         'bias': 0.01,
     }
 
+    # NOTE: Add more in the future as needed.
+    SUPPORTED_HETERO_MODELS = [
+        HGTConv,
+        HANConv,
+        HeteroConv,
+    ]
+
     def __init__(self, epochs: int, lr: float = 0.003, **kwargs):
         super().__init__()
         self.epochs = epochs
@@ -75,11 +87,13 @@ class PGExplainer(ExplainerAlgorithm):
         )
         self.optimizer = torch.optim.Adam(self.mlp.parameters(), lr=lr)
         self._curr_epoch = -1
+        self.is_hetero = False
 
     def reset_parameters(self):
         r"""Resets all learnable parameters of the module."""
         reset(self.mlp)
 
+    @overload
     def train(
         self,
         epoch: int,
@@ -90,17 +104,44 @@ class PGExplainer(ExplainerAlgorithm):
         target: Tensor,
         index: Optional[Union[int, Tensor]] = None,
         **kwargs,
-    ):
+    ) -> float:
+        ...
+
+    @overload
+    def train(
+        self,
+        epoch: int,
+        model: torch.nn.Module,
+        x: Dict[NodeType, Tensor],
+        edge_index: Dict[EdgeType, Tensor],
+        *,
+        target: Tensor,
+        index: Optional[Union[int, Tensor]] = None,
+        **kwargs,
+    ) -> float:
+        ...
+
+    def train(
+        self,
+        epoch: int,
+        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,
+    ) -> float:
         r"""Trains the underlying explainer model.
         Needs to be called before being able to make predictions.
 
         Args:
             epoch (int): The current epoch of the training phase.
             model (torch.nn.Module): The model to explain.
-            x (torch.Tensor): The input node features of a
-                homogeneous graph.
-            edge_index (torch.Tensor): The input edge indices of a homogeneous
-                graph.
+            x (torch.Tensor or Dict[str, torch.Tensor]): The input node
+                features. Can be either homogeneous or heterogeneous.
+            edge_index (torch.Tensor or Dict[Tuple[str, str, str]): The input
+                edge indices. Can be either homogeneous or heterogeneous.
             target (torch.Tensor): The target of the model.
             index (int or torch.Tensor, optional): The index of the model
                 output to explain. Needs to be a single index.
@@ -108,9 +149,9 @@ class PGExplainer(ExplainerAlgorithm):
             **kwargs (optional): Additional keyword arguments passed to
                 :obj:`model`.
         """
-        if isinstance(x, dict) or isinstance(edge_index, dict):
-            raise ValueError(f"Heterogeneous graphs not yet supported in "
-                             f"'{self.__class__.__name__}'")
+        self.is_hetero = isinstance(x, dict)
+        if self.is_hetero:
+            assert isinstance(edge_index, dict)
 
         if self.model_config.task_level == ModelTaskLevel.node:
             if index is None:
@@ -121,35 +162,68 @@ class PGExplainer(ExplainerAlgorithm):
                 raise ValueError(f"Only scalars are supported for the 'index' "
                                  f"argument in '{self.__class__.__name__}'")
 
-        z = get_embeddings(model, x, edge_index, **kwargs)[-1]
+        # Get embeddings based on whether the graph is homogeneous or
+        # heterogeneous
+        node_embeddings = self._get_embeddings(model, x, edge_index, **kwargs)
 
+        # Train the model
         self.optimizer.zero_grad()
         temperature = self._get_temperature(epoch)
 
-        inputs = self._get_inputs(z, edge_index, index)
-        logits = self.mlp(inputs).view(-1)
-        edge_mask = self._concrete_sample(logits, temperature)
-        set_masks(model, edge_mask, edge_index, apply_sigmoid=True)
-
-        if self.model_config.task_level == ModelTaskLevel.node:
-            _, hard_edge_mask = self._get_hard_masks(model, index, edge_index,
-                                                     num_nodes=x.size(0))
-            edge_mask = edge_mask[hard_edge_mask]
-
+        # Process embeddings and generate edge masks
+        edge_mask = self._generate_edge_masks(node_embeddings, edge_index,
+                                              index, temperature)
+
+        # Apply masks to the model
+        if self.is_hetero:
+            set_hetero_masks(model, edge_mask, edge_index, apply_sigmoid=True)
+
+            # For node-level tasks, we can compute hard masks
+            if self.model_config.task_level == ModelTaskLevel.node:
+                # Process each edge type separately
+                for edge_type, mask in edge_mask.items():
+                    # Get the edge indices for this edge type
+                    edges = edge_index[edge_type]
+                    src_type, _, dst_type = edge_type
+
+                    # Get hard masks for this specific edge type
+                    _, hard_mask = self._get_hard_masks(
+                        model, index, edges,
+                        num_nodes=max(x[src_type].size(0),
+                                      x[dst_type].size(0)))
+
+                    edge_mask[edge_type] = mask[hard_mask]
+        else:
+            # Apply masks for homogeneous graphs
+            set_masks(model, edge_mask, edge_index, apply_sigmoid=True)
+
+            # For node-level tasks, we may need to apply hard masks
+            hard_edge_mask = None
+            if self.model_config.task_level == ModelTaskLevel.node:
+                _, hard_edge_mask = self._get_hard_masks(
+                    model, index, edge_index, num_nodes=x.size(0))
+                edge_mask = edge_mask[hard_edge_mask]
+
+        # Forward pass with masks applied
         y_hat, y = model(x, edge_index, **kwargs), target
 
         if index is not None:
             y_hat, y = y_hat[index], y[index]
 
+        # Calculate loss
         loss = self._loss(y_hat, y, edge_mask)
+
+        # Backward pass and optimization
         loss.backward()
         self.optimizer.step()
 
+        # Clean up
         clear_masks(model)
         self._curr_epoch = epoch
 
         return float(loss)
 
+    @overload
     def forward(
         self,
         model: torch.nn.Module,
@@ -160,9 +234,32 @@ class PGExplainer(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__}'")
+        ...
+
+    @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)
 
         if self._curr_epoch < self.epochs - 1:  # Safety check:
             raise ValueError(f"'{self.__class__.__name__}' is not yet fully "
@@ -171,7 +268,6 @@ class PGExplainer(ExplainerAlgorithm):
                              f"the underlying explainer model by running "
                              f"`explainer.algorithm.train(...)`.")
 
-        hard_edge_mask = None
         if self.model_config.task_level == ModelTaskLevel.node:
             if index is None:
                 raise ValueError(f"The 'index' argument needs to be provided "
@@ -181,20 +277,55 @@ class PGExplainer(ExplainerAlgorithm):
                 raise ValueError(f"Only scalars are supported for the 'index' "
                                  f"argument in '{self.__class__.__name__}'")
 
-            # We need to compute hard masks to properly clean up edges and
-            # nodes attributions not involved during message passing:
-            _, hard_edge_mask = self._get_hard_masks(model, index, edge_index,
-                                                     num_nodes=x.size(0))
+        # Get embeddings
+        node_embeddings = self._get_embeddings(model, x, edge_index, **kwargs)
+
+        # Generate explanations
+        if self.is_hetero:
+            # Generate edge masks for each edge type
+            edge_masks = {}
+
+            # Generate masks for each edge type
+            for edge_type, edge_idx in edge_index.items():
+                src_node_type, _, dst_node_type = edge_type
+
+                assert src_node_type in node_embeddings
+                assert dst_node_type in node_embeddings
+
+                inputs = self._get_inputs_hetero(node_embeddings, edge_type,
+                                                 edge_idx, index)
+                logits = self.mlp(inputs).view(-1)
+
+                # For node-level explanations, get hard masks for this
+                # specific edge type
+                hard_edge_mask = None
+                if self.model_config.task_level == ModelTaskLevel.node:
+                    _, hard_edge_mask = self._get_hard_masks(
+                        model, index, edge_idx,
+                        num_nodes=max(x[src_node_type].size(0),
+                                      x[dst_node_type].size(0)))
 
-        z = get_embeddings(model, x, edge_index, **kwargs)[-1]
+                # Apply hard mask if available and it has any True values
+                edge_masks[edge_type] = self._post_process_mask(
+                    logits, hard_edge_mask, apply_sigmoid=True)
 
-        inputs = self._get_inputs(z, edge_index, index)
-        logits = self.mlp(inputs).view(-1)
+            explanation = HeteroExplanation()
+            explanation.set_value_dict('edge_mask', edge_masks)
+            return explanation
+        else:
+            hard_edge_mask = None
+            if self.model_config.task_level == ModelTaskLevel.node:
+                # We need to compute hard masks to properly clean up edges
+                _, hard_edge_mask = self._get_hard_masks(
+                    model, index, edge_index, num_nodes=x.size(0))
 
-        edge_mask = self._post_process_mask(logits, hard_edge_mask,
-                                            apply_sigmoid=True)
+            inputs = self._get_inputs(node_embeddings, edge_index, index)
+            logits = self.mlp(inputs).view(-1)
 
-        return Explanation(edge_mask=edge_mask)
+            edge_mask = self._post_process_mask(logits, hard_edge_mask,
+                                                apply_sigmoid=True)
+
+            return Explanation(edge_mask=edge_mask)
 
     def supports(self) -> bool:
         explanation_type = self.explainer_config.explanation_type
@@ -222,6 +353,76 @@ class PGExplainer(ExplainerAlgorithm):
 
     ###########################################################################
 
+    def _get_embeddings(self, model: torch.nn.Module, x: Union[Tensor,
+                                                               Dict[NodeType,
+                                                                    Tensor]],
+                        edge_index: Union[Tensor, Dict[EdgeType, Tensor]],
+                        **kwargs) -> Union[Tensor, Dict[NodeType, Tensor]]:
+        """Get embeddings from the model based on input type."""
+        if self.is_hetero:
+            # For heterogeneous graphs, get embeddings for each node type
+            embeddings_dict = get_embeddings_hetero(
+                model,
+                self.SUPPORTED_HETERO_MODELS,
+                x,
+                edge_index,
+                **kwargs,
+            )
+
+            # Use the last layer's embeddings for each node type
+            last_embedding_dict = {
+                node_type: embs[-1] if embs and len(embs) > 0 else None
+                for node_type, embs in embeddings_dict.items()
+            }
+
+            # Skip if no embeddings were captured
+            if not any(emb is not None
+                       for emb in last_embedding_dict.values()):
+                raise ValueError(
+                    "No embeddings were captured from the model. "
+                    "Please check if the model architecture is supported.")
+
+            return last_embedding_dict
+        else:
+            # For homogeneous graphs, get embeddings directly
+            return get_embeddings(model, x, edge_index, **kwargs)[-1]
+
+    def _generate_edge_masks(
+            self, emb: Union[Tensor, Dict[NodeType, Tensor]],
+            edge_index: Union[Tensor,
+                              Dict[EdgeType,
+                                   Tensor]], index: Optional[Union[int,
+                                                                   Tensor]],
+            temperature: float) -> Union[Tensor, Dict[EdgeType, Tensor]]:
+        """Generate edge masks based on embeddings."""
+        if self.is_hetero:
+            # For heterogeneous graphs, generate masks for each edge type
+            edge_masks = {}
+
+            for edge_type, edge_idx in edge_index.items():
+                src, _, dst = edge_type
+
+                assert src in emb and dst in emb
+                # Generate inputs for this edge type
+                inputs = self._get_inputs_hetero(emb, edge_type, edge_idx,
+                                                 index)
+                logits = self.mlp(inputs).view(-1)
+                edge_masks[edge_type] = self._concrete_sample(
+                    logits, temperature)
+
+            # Ensure we have at least one valid edge mask
+            if not edge_masks:
+                raise ValueError(
+                    "Could not generate edge masks for any edge type. "
+                    "Please ensure the model architecture is supported.")
+
+            return edge_masks
+        else:
+            # For homogeneous graphs, generate a single mask
+            inputs = self._get_inputs(emb, edge_index, index)
+            logits = self.mlp(inputs).view(-1)
+            return self._concrete_sample(logits, temperature)
+
     def _get_inputs(self, embedding: Tensor, edge_index: Tensor,
                     index: Optional[int] = None) -> Tensor:
         zs = [embedding[edge_index[0]], embedding[edge_index[1]]]
@@ -230,6 +431,27 @@ class PGExplainer(ExplainerAlgorithm):
             zs.append(embedding[index].view(1, -1).repeat(zs[0].size(0), 1))
         return torch.cat(zs, dim=-1)
 
+    def _get_inputs_hetero(self, embedding_dict: Dict[NodeType, Tensor],
+                           edge_type: Tuple[str, str, str], edge_index: Tensor,
+                           index: Optional[int] = None) -> Tensor:
+        src, _, dst = edge_type
+
+        # Get embeddings for source and destination nodes
+        src_emb = embedding_dict[src]
+        dst_emb = embedding_dict[dst]
+
+        # Source and destination node embeddings
+        zs = [src_emb[edge_index[0]], dst_emb[edge_index[1]]]
+
+        # For node-level explanations, add the target node embedding
+        if self.model_config.task_level == ModelTaskLevel.node:
+            assert index is not None
+            # Assuming index refers to a node of type 'src'
+            target_emb = src_emb[index].view(1, -1).repeat(zs[0].size(0), 1)
+            zs.append(target_emb)
+
+        return torch.cat(zs, dim=-1)
+
     def _get_temperature(self, epoch: int) -> float:
         temp = self.coeffs['temp']
         return temp[0] * pow(temp[1] / temp[0], epoch / self.epochs)
@@ -240,19 +462,55 @@ class PGExplainer(ExplainerAlgorithm):
         eps = (1 - 2 * bias) * torch.rand_like(logits) + bias
         return (eps.log() - (1 - eps).log() + logits) / temperature
 
-    def _loss(self, y_hat: Tensor, y: Tensor, edge_mask: Tensor) -> Tensor:
+    def _loss(self, y_hat: Tensor, y: Tensor,
+              edge_mask: Union[Tensor, Dict[EdgeType, 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:
+            loss = self._apply_hetero_regularization(loss, edge_mask)
+        else:
+            loss = self._apply_homo_regularization(loss, edge_mask)
+
+        return loss
+
+    def _calculate_base_loss(self, y_hat: Tensor, y: Tensor) -> Tensor:
+        """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)
-
-        # Regularization loss:
-        mask = edge_mask.sigmoid()
+            return self._loss_regression(y_hat, y)
+        else:
+            raise ValueError(
+                f"Unsupported model mode: {self.model_config.mode}")
+
+    def _apply_hetero_regularization(
+            self, loss: Tensor, edge_mask: Dict[EdgeType, Tensor]) -> Tensor:
+        """Apply regularization for heterogeneous graph."""
+        for _, mask in edge_mask.items():
+            loss = self._add_mask_regularization(loss, mask)
+
+        return loss
+
+    def _apply_homo_regularization(self, loss: Tensor,
+                                   edge_mask: Tensor) -> Tensor:
+        """Apply regularization for homogeneous graph."""
+        return self._add_mask_regularization(loss, edge_mask)
+
+    def _add_mask_regularization(self, loss: Tensor, mask: Tensor) -> Tensor:
+        """Add size and entropy regularization for a mask."""
+        # Apply sigmoid for mask values
+        mask = mask.sigmoid()
+
+        # Size regularization
         size_loss = mask.sum() * self.coeffs['edge_size']
-        mask = 0.99 * mask + 0.005
-        mask_ent = -mask * mask.log() - (1 - mask) * (1 - mask).log()
+
+        # Entropy regularization
+        masked = 0.99 * mask + 0.005
+        mask_ent = -masked * masked.log() - (1 - masked) * (1 - masked).log()
         mask_ent_loss = mask_ent.mean() * self.coeffs['edge_ent']
 
         return loss + size_loss + mask_ent_loss

torch_geometric/utils/__init__.py

@@ -53,7 +53,7 @@ from ._negative_sampling import (negative_sampling, batched_negative_sampling,
                                  structured_negative_sampling_feasible)
 from .augmentation import shuffle_node, mask_feature, add_random_edge
 from ._tree_decomposition import tree_decomposition
-from .embedding import get_embeddings
+from .embedding import get_embeddings, get_embeddings_hetero
 from ._trim_to_layer import trim_to_layer
 from .ppr import get_ppr
 from ._train_test_split_edges import train_test_split_edges
@@ -145,6 +145,7 @@ __all__ = [
     'add_random_edge',
     'tree_decomposition',
     'get_embeddings',
+    'get_embeddings_hetero',
     'trim_to_layer',
     'get_ppr',
     'train_test_split_edges',

torch_geometric/utils/embedding.py

@@ -1,9 +1,11 @@
 import warnings
-from typing import Any, List
+from typing import Any, Dict, List, Optional, Type
 
 import torch
 from torch import Tensor
 
+from torch_geometric.typing import NodeType
+
 
 def get_embeddings(
     model: torch.nn.Module,
@@ -52,3 +54,88 @@ def get_embeddings(
         handle.remove()
 
     return embeddings
+
+
+def get_embeddings_hetero(
+    model: torch.nn.Module,
+    supported_models: Optional[List[Type[torch.nn.Module]]] = None,
+    *args: Any,
+    **kwargs: Any,
+) -> Dict[NodeType, List[Tensor]]:
+    """Returns the output embeddings of all
+    :class:`~torch_geometric.nn.conv.MessagePassing` layers in a heterogeneous
+    :obj:`model`, organized by edge type.
+
+    Internally, this method registers forward hooks on all modules that process
+    heterogeneous graphs in the model and runs the forward pass of the model.
+    For heterogeneous models, the output is a dictionary where each key is a
+    node type and each value is a list of embeddings from different layers.
+
+    Args:
+        model (torch.nn.Module): The heterogeneous GNN model.
+        supported_models (List[Type[torch.nn.Module]], optional): A list of
+            supported model classes. If not provided, defaults to
+            [HGTConv, HANConv, HeteroConv].
+        *args: Arguments passed to the model.
+        **kwargs (optional): Additional keyword arguments passed to the model.
+
+    Returns:
+        Dict[NodeType, List[Tensor]]: A dictionary mapping each node type to
+        a list of embeddings from different layers.
+    """
+    from torch_geometric.nn import HANConv, HeteroConv, HGTConv
+    if not supported_models:
+        supported_models = [HGTConv, HANConv, HeteroConv]
+
+    # Dictionary to store node embeddings by type
+    node_embeddings_dict: Dict[NodeType, List[Tensor]] = {}
+
+    # Hook function to capture node embeddings
+    def hook(model: torch.nn.Module, inputs: Any, outputs: Any) -> None:
+        # Check if the outputs is a dictionary mapping node types to embeddings
+        if isinstance(outputs, dict) and outputs:
+            # Store embeddings for each node type
+            for node_type, embedding in outputs.items():
+                # Made sure that the outputs are a dictionary mapping node
+                # types to embeddings and remove the false positives.
+                if node_type not in node_embeddings_dict:
+                    node_embeddings_dict[node_type] = []
+                node_embeddings_dict[node_type].append(embedding.clone())
+
+    # List to store hook handles
+    hook_handles = []
+
+    # Find ModuleDict objects in the model
+    for _, module in model.named_modules():
+        # Handle the native heterogenous models, e.g. HGTConv, HANConv
+        # and HeteroConv, etc.
+        if isinstance(module, tuple(supported_models)):
+            hook_handles.append(module.register_forward_hook(hook))
+        else:
+            # Handle the heterogenous models that are generated by calling
+            # to_hetero() on the homogeneous models.
+            submodules = list(module.children())
+            submodules_contains_module_dict = any([
+                isinstance(submodule, torch.nn.ModuleDict)
+                for submodule in submodules
+            ])
+            if submodules_contains_module_dict:
+                hook_handles.append(module.register_forward_hook(hook))
+
+    if len(hook_handles) == 0:
+        warnings.warn("The 'model' does not have any heterogenous "
+                      "'MessagePassing' layers")
+
+    # Run the model forward pass
+    training = model.training
+    model.eval()
+
+    with torch.no_grad():
+        model(*args, **kwargs)
+    model.train(training)
+
+    # Clean up hooks
+    for handle in hook_handles:
+        handle.remove()
+
+    return node_embeddings_dict