pyg-nightly 2.7.0.dev20241009__py3-none-any.whl → 2.8.0.dev20251228__py3-none-any.whl

torch_geometric/edge_index.py

@@ -17,6 +17,7 @@ from typing import (
     overload,
 )
 
+import numpy as np
 import torch
 import torch.utils._pytree as pytree
 from torch import Tensor
@@ -183,7 +184,7 @@ class EdgeIndex(Tensor):
 
         edge_index = EdgeIndex(
             [[0, 1, 1, 2],
-             [1, 0, 2, 1]]
+             [1, 0, 2, 1]],
             sparse_size=(3, 3),
             sort_order='row',
             is_undirected=True,
@@ -210,7 +211,7 @@ class EdgeIndex(Tensor):
         assert not edge_index.is_undirected
 
         # Sparse-Dense Matrix Multiplication:
-        out = edge_index.flip(0) @ torch.randn(3, 16)
+        out = edge_index.flip(0) @ torch.randn(3, 16)
         assert out.size() == (3, 16)
     """
     # See "https://pytorch.org/docs/stable/notes/extending.html"
@@ -297,8 +298,7 @@ class EdgeIndex(Tensor):
                 indptr = None
             data = torch.stack([row, col], dim=0)
 
-        if (torch_geometric.typing.WITH_PT112
-                and data.layout == torch.sparse_csc):
+        if data.layout == torch.sparse_csc:
             row = data.row_indices()
             indptr = data.ccol_indices()
 
@@ -325,7 +325,7 @@ class EdgeIndex(Tensor):
             elif sparse_size[0] is None and sparse_size[1] is not None:
                 sparse_size = (sparse_size[1], sparse_size[1])
 
-        out = Tensor._make_wrapper_subclass(  # type: ignore
+        out = Tensor._make_wrapper_subclass(
             cls,
             size=data.size(),
             strides=data.stride(),
@@ -803,7 +803,7 @@ class EdgeIndex(Tensor):
 
         size = self.get_sparse_size()
         if value is not None and value.dim() > 1:
-            size = size + value.size()[1:]  # type: ignore
+            size = size + value.size()[1:]
 
         out = torch.full(size, fill_value, dtype=dtype, device=self.device)
         out[self._data[0], self._data[1]] = value if value is not None else 1
@@ -881,10 +881,6 @@ class EdgeIndex(Tensor):
                 If not specified, non-zero elements will be assigned a value of
                 :obj:`1.0`. (default: :obj:`None`)
         """
-        if not torch_geometric.typing.WITH_PT112:
-            raise NotImplementedError(
-                "'to_sparse_csc' not supported for PyTorch < 1.12")
-
         (colptr, row), perm = self.get_csc()
         if value is not None and perm is not None:
             value = value[perm]
@@ -921,7 +917,7 @@ class EdgeIndex(Tensor):
             return self.to_sparse_coo(value)
         if layout == torch.sparse_csr:
             return self.to_sparse_csr(value)
-        if torch_geometric.typing.WITH_PT112 and layout == torch.sparse_csc:
+        if layout == torch.sparse_csc:
             return self.to_sparse_csc(value)
 
         raise ValueError(f"Unexpected tensor layout (got '{layout}')")
@@ -1190,10 +1186,10 @@ class EdgeIndex(Tensor):
         return edge_index
 
     # Prevent auto-wrapping outputs back into the proper subclass type:
-    __torch_function__ = torch._C._disabled_torch_function_impl
+    __torch_function__ = torch._C._disabled_torch_function_impl  # type: ignore
 
     @classmethod
-    def __torch_dispatch__(
+    def __torch_dispatch__(  # type: ignore
         cls: Type,
         func: Callable[..., Any],
         types: Iterable[Type[Any]],
@@ -1246,6 +1242,14 @@ class EdgeIndex(Tensor):
         return torch._tensor_str._add_suffixes(prefix + tensor_str, suffixes,
                                                indent, force_newline=False)
 
+    def tolist(self) -> List[Any]:
+        """"""  # noqa: D419
+        return self._data.tolist()
+
+    def numpy(self, *, force: bool = False) -> np.ndarray:
+        """"""  # noqa: D419
+        return self._data.numpy(force=force)
+
     # Helpers #################################################################
 
     def _shallow_copy(self) -> 'EdgeIndex':
@@ -1478,7 +1482,7 @@ def _slice(
     step: int = 1,
 ) -> Union[EdgeIndex, Tensor]:
 
-    if ((start is None or start <= 0)
+    if ((start is None or start == 0 or start <= -input.size(dim))
             and (end is None or end > input.size(dim)) and step == 1):
         return input._shallow_copy()  # No-op.
 

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:

torch_geometric/explain/algorithm/base.py

@@ -166,7 +166,7 @@ class ExplainerAlgorithm(torch.nn.Module):
         elif self.model_config.return_type == ModelReturnType.probs:
             loss_fn = F.binary_cross_entropy
         else:
-            assert False
+            raise AssertionError()
 
         return loss_fn(y_hat.view_as(y), y.float())
 
@@ -183,7 +183,7 @@ class ExplainerAlgorithm(torch.nn.Module):
         elif self.model_config.return_type == ModelReturnType.log_probs:
             loss_fn = F.nll_loss
         else:
-            assert False
+            raise AssertionError()
 
         return loss_fn(y_hat, y)
 

torch_geometric/explain/algorithm/captum.py

@@ -190,7 +190,7 @@ def to_captum_input(
 
     Args:
         x (torch.Tensor or Dict[NodeType, torch.Tensor]): The node features.
-            For heterogeneous graphs this is a dictionary holding node featues
+            For heterogeneous graphs this is a dictionary holding node features
             for each node type.
         edge_index(torch.Tensor or Dict[EdgeType, torch.Tensor]): The edge
             indices. For heterogeneous graphs this is a dictionary holding the

torch_geometric/explain/algorithm/captum_explainer.py

@@ -73,7 +73,8 @@ class CaptumExplainer(ExplainerAlgorithm):
                              f"{self.attribution_method_class.__name__}")
 
         if kwargs.get('internal_batch_size', 1) != 1:
-            warnings.warn("Overriding 'internal_batch_size' to 1")
+            warnings.warn("Overriding 'internal_batch_size' to 1",
+                          stacklevel=2)
 
         if 'internal_batch_size' in self._get_attribute_parameters():
             kwargs['internal_batch_size'] = 1