pyg-nightly 2.7.0.dev20250423__py3-none-any.whl → 2.7.0.dev20250424__py3-none-any.whl

{pyg_nightly-2.7.0.dev20250423.dist-info → pyg_nightly-2.7.0.dev20250424.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: pyg-nightly
-Version: 2.7.0.dev20250423
+Version: 2.7.0.dev20250424
 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.dev20250423.dist-info → pyg_nightly-2.7.0.dev20250424.dist-info}/RECORD

@@ -1,4 +1,4 @@
-torch_geometric/__init__.py,sha256=BRFGlcJ1nDb4azLG0EB7y7xon3yTqDiN8YOQ9ACnscI,1978
+torch_geometric/__init__.py,sha256=ERKw5z1mre0dUq9Ds7it-e7OfZSTr-CjZED7ofIMMW4,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
@@ -195,7 +195,7 @@ torch_geometric/explain/config.py,sha256=_0j67NAwPwjrWHPncNywCT-oKyMiryJNxufxVN1
 torch_geometric/explain/explainer.py,sha256=8_NZTmlT4WO9RgKxpSUQRt3rbVwFURF5bSWOPlfOLjA,10667
 torch_geometric/explain/explanation.py,sha256=Z2NlgavEnq0QadEr6p6pxAhV6lU7WrlcJLFWbTdsmvg,14903
 torch_geometric/explain/algorithm/__init__.py,sha256=fE29xbd0bPxg-EfrB2BDmmY9QnyO-7TgvYduGHofm5o,496
-torch_geometric/explain/algorithm/attention_explainer.py,sha256=iRWgrUVoAn42DpVPE0jZclLB6OtUOArKl5dn53WmCc4,4545
+torch_geometric/explain/algorithm/attention_explainer.py,sha256=65iGLmOt00ERtBDVxAoydIchykdWZU24aXzSzUGzQEI,11304
 torch_geometric/explain/algorithm/base.py,sha256=wwJcREUFKDLFUDjRa9o4X3DWqQgMvhS3Iciwb6Evtjc,6922
 torch_geometric/explain/algorithm/captum.py,sha256=k6hNgC5Kn9lVirOYVJzej8-hRuf5C2mPFUXFLd2wWsY,12857
 torch_geometric/explain/algorithm/captum_explainer.py,sha256=oz-c40hvdzii4_chEQPHzQo_dFjHr9HLuJhDLsqRIVU,7346
@@ -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.dev20250423.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
-pyg_nightly-2.7.0.dev20250423.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
-pyg_nightly-2.7.0.dev20250423.dist-info/METADATA,sha256=OOH3EHfpXcmpGeyXvgYej5Brib1sYdXOJDrJkPwGaXs,62979
-pyg_nightly-2.7.0.dev20250423.dist-info/RECORD,,
+pyg_nightly-2.7.0.dev20250424.dist-info/licenses/LICENSE,sha256=ic-27cMJc1kWoMEYncz3Ya3Ur2Bi3bNLWib2DT763-o,1067
+pyg_nightly-2.7.0.dev20250424.dist-info/WHEEL,sha256=G2gURzTEtmeR8nrdXUJfNiB3VYVxigPQ-bEQujpNiNs,82
+pyg_nightly-2.7.0.dev20250424.dist-info/METADATA,sha256=fQqhKZ4bkgt-_tyrSNhIMz0ZUixa4ctPYvIPGdWCQzo,62979
+pyg_nightly-2.7.0.dev20250424.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.dev20250423'
+__version__ = '2.7.0.dev20250424'
 
 __all__ = [
     'Index',

torch_geometric/explain/algorithm/attention_explainer.py

@@ -1,13 +1,14 @@
 import logging
-from typing import List, Optional, Union
+from typing import Dict, List, Optional, Union, overload
 
 import torch
 from torch import Tensor
 
-from torch_geometric.explain import Explanation
+from torch_geometric.explain import Explanation, HeteroExplanation
 from torch_geometric.explain.algorithm import ExplainerAlgorithm
 from torch_geometric.explain.config import ExplanationType, ModelTaskLevel
 from torch_geometric.nn.conv.message_passing import MessagePassing
+from torch_geometric.typing import EdgeType, NodeType
 
 
 class AttentionExplainer(ExplainerAlgorithm):
@@ -26,7 +27,9 @@ class AttentionExplainer(ExplainerAlgorithm):
     def __init__(self, reduce: str = 'max'):
         super().__init__()
         self.reduce = reduce
+        self.is_hetero = False
 
+    @overload
     def forward(
         self,
         model: torch.nn.Module,
@@ -37,65 +40,252 @@ class AttentionExplainer(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__}'")
+        ...
 
-        hard_edge_mask = None
-        if self.model_config.task_level == ModelTaskLevel.node:
-            # We need to compute the hard edge mask to properly clean up edge
-            # attributions not involved during message passing:
-            _, hard_edge_mask = self._get_hard_masks(model, index, edge_index,
-                                                     num_nodes=x.size(0))
+    @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]:
+        """Generate explanations based on attention coefficients."""
+        self.is_hetero = isinstance(x, dict)
+
+        # Collect attention coefficients
+        alphas_dict = self._collect_attention_coefficients(
+            model, x, edge_index, **kwargs)
+
+        # Process attention coefficients
+        if self.is_hetero:
+            return self._create_hetero_explanation(model, alphas_dict,
+                                                   edge_index, index, x)
+        else:
+            return self._create_homo_explanation(model, alphas_dict,
+                                                 edge_index, index, x)
+
+    @overload
+    def _collect_attention_coefficients(
+        self,
+        model: torch.nn.Module,
+        x: Tensor,
+        edge_index: Tensor,
+        **kwargs,
+    ) -> List[Tensor]:
+        ...
+
+    @overload
+    def _collect_attention_coefficients(
+        self,
+        model: torch.nn.Module,
+        x: Dict[NodeType, Tensor],
+        edge_index: Dict[EdgeType, Tensor],
+        **kwargs,
+    ) -> Dict[EdgeType, List[Tensor]]:
+        ...
+
+    def _collect_attention_coefficients(
+        self,
+        model: torch.nn.Module,
+        x: Union[Tensor, Dict[NodeType, Tensor]],
+        edge_index: Union[Tensor, Dict[EdgeType, Tensor]],
+        **kwargs,
+    ) -> Union[List[Tensor], Dict[EdgeType, List[Tensor]]]:
+        """Collect attention coefficients from model layers."""
+        if self.is_hetero:
+            # For heterogeneous graphs, store alphas by edge type
+            alphas_dict: Dict[EdgeType, List[Tensor]] = {}
+
+            # Get list of edge types
+            edge_types = list(edge_index.keys())
+
+            # Hook function to capture attention coefficients by edge type
+            def hook(module, msg_kwargs, out):
+                # Find edge type from the module's full name
+                module_name = getattr(module, '_name', None)
+                if module_name is None:
+                    return
 
-        alphas: List[Tensor] = []
+                edge_type = None
+                for edge_tuple in edge_types:
+                    src_type, edge_name, dst_type = edge_tuple
+                    # Check if all components appear in the module name in
+                    # order
+                    try:
+                        src_idx = module_name.index(src_type)
+                        edge_idx = module_name.index(edge_name, src_idx)
+                        dst_idx = module_name.index(dst_type, edge_idx)
+                        if src_idx < edge_idx < dst_idx:
+                            edge_type = edge_tuple
+                            break
+                    except ValueError:  # Component not found
+                        continue
+
+                if edge_type is None:
+                    return
+
+                if edge_type not in alphas_dict:
+                    alphas_dict[edge_type] = []
+
+                # Extract alpha from message kwargs or module
+                if 'alpha' in msg_kwargs[0]:
+                    alphas_dict[edge_type].append(
+                        msg_kwargs[0]['alpha'].detach())
+                elif getattr(module, '_alpha', None) is not None:
+                    alphas_dict[edge_type].append(module._alpha.detach())
+        else:
+            # For homogeneous graphs, store all alphas in a list
+            alphas: List[Tensor] = []
 
-        def hook(module, msg_kwargs, out):
-            if 'alpha' in msg_kwargs[0]:
-                alphas.append(msg_kwargs[0]['alpha'].detach())
-            elif getattr(module, '_alpha', None) is not None:
-                alphas.append(module._alpha.detach())
+            def hook(module, msg_kwargs, out):
+                if 'alpha' in msg_kwargs[0]:
+                    alphas.append(msg_kwargs[0]['alpha'].detach())
+                elif getattr(module, '_alpha', None) is not None:
+                    alphas.append(module._alpha.detach())
 
+        # Register hooks for all message passing modules
         hook_handles = []
-        for module in model.modules():  # Register message forward hooks:
-            if (isinstance(module, MessagePassing)
-                    and module.explain is not False):
+        for name, module in model.named_modules():
+            if isinstance(module,
+                          MessagePassing) and module.explain is not False:
+                # Store name for hetero graph lookup in the hook
+                if self.is_hetero:
+                    module._name = name
+
                 hook_handles.append(module.register_message_forward_hook(hook))
 
+        # Forward pass to collect attention coefficients.
         model(x, edge_index, **kwargs)
 
-        for handle in hook_handles:  # Remove hooks:
+        # Remove hooks
+        for handle in hook_handles:
             handle.remove()
 
-        if len(alphas) == 0:
-            raise ValueError("Could not collect any attention coefficients. "
-                             "Please ensure that your model is using "
-                             "attention-based GNN layers.")
+        # Check if we collected any attention coefficients.
+        if self.is_hetero:
+            if not alphas_dict:
+                raise ValueError(
+                    "Could not collect any attention coefficients. "
+                    "Please ensure that your model is using "
+                    "attention-based GNN layers.")
+            return alphas_dict
+        else:
+            if not alphas:
+                raise ValueError(
+                    "Could not collect any attention coefficients. "
+                    "Please ensure that your model is using "
+                    "attention-based GNN layers.")
+            return alphas
 
+    def _process_attention_coefficients(
+        self,
+        alphas: List[Tensor],
+        edge_index_size: int,
+    ) -> Tensor:
+        """Process collected attention coefficients into a single mask."""
         for i, alpha in enumerate(alphas):
-            alpha = alpha[:edge_index.size(1)]  # Respect potential self-loops.
+            # Ensure alpha doesn't exceed edge_index size
+            alpha = alpha[:edge_index_size]
+
+            # Reduce multi-head attention
             if alpha.dim() == 2:
                 alpha = getattr(torch, self.reduce)(alpha, dim=-1)
-                if isinstance(alpha, tuple):  # Respect `torch.max`:
+                if isinstance(alpha, tuple):  # Handle torch.max output
                     alpha = alpha[0]
             elif alpha.dim() > 2:
-                raise ValueError(f"Can not reduce attention coefficients of "
+                raise ValueError(f"Cannot reduce attention coefficients of "
                                  f"shape {list(alpha.size())}")
             alphas[i] = alpha
 
+        # Combine attention coefficients across layers
         if len(alphas) > 1:
             alpha = torch.stack(alphas, dim=-1)
             alpha = getattr(torch, self.reduce)(alpha, dim=-1)
-            if isinstance(alpha, tuple):  # Respect `torch.max`:
+            if isinstance(alpha, tuple):  # Handle torch.max output
                 alpha = alpha[0]
         else:
             alpha = alphas[0]
 
+        return alpha
+
+    def _create_homo_explanation(
+        self,
+        model: torch.nn.Module,
+        alphas: List[Tensor],
+        edge_index: Tensor,
+        index: Optional[Union[int, Tensor]],
+        x: Tensor,
+    ) -> Explanation:
+        """Create explanation for homogeneous graph."""
+        # Get hard edge mask for node-level tasks
+        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))
+
+        # Process attention coefficients
+        alpha = self._process_attention_coefficients(alphas,
+                                                     edge_index.size(1))
+
+        # Post-process mask with hard edge mask if needed
         alpha = self._post_process_mask(alpha, hard_edge_mask,
                                         apply_sigmoid=False)
 
         return Explanation(edge_mask=alpha)
 
+    def _create_hetero_explanation(
+        self,
+        model: torch.nn.Module,
+        alphas_dict: Dict[EdgeType, List[Tensor]],
+        edge_index: Dict[EdgeType, Tensor],
+        index: Optional[Union[int, Tensor]],
+        x: Dict[NodeType, Tensor],
+    ) -> HeteroExplanation:
+        """Create explanation for heterogeneous graph."""
+        edge_masks_dict = {}
+
+        # Process each edge type separately
+        for edge_type, alphas in alphas_dict.items():
+            if not alphas:
+                continue
+
+            # Get hard edge mask for node-level tasks
+            hard_edge_mask = None
+            if self.model_config.task_level == ModelTaskLevel.node:
+                src_type, _, dst_type = edge_type
+                _, hard_edge_mask = self._get_hard_masks(
+                    model, index, edge_index[edge_type],
+                    num_nodes=max(x[src_type].size(0), x[dst_type].size(0)))
+
+            # Process attention coefficients for this edge type
+            alpha = self._process_attention_coefficients(
+                alphas, edge_index[edge_type].size(1))
+
+            # Apply hard mask if available
+            edge_masks_dict[edge_type] = self._post_process_mask(
+                alpha, hard_edge_mask, apply_sigmoid=False)
+
+        # Create heterogeneous explanation
+        explanation = HeteroExplanation()
+        explanation.set_value_dict('edge_mask', edge_masks_dict)
+        return explanation
+
     def supports(self) -> bool:
         explanation_type = self.explainer_config.explanation_type
         if explanation_type != ExplanationType.model: