Repository: pyg-team/pytorch_geometric
Pull request #10169 (2025-04-06 00:28:28)

Added support for heterogenous graphs in `AttentionExplainer`
Fixes https://github.com/pyg-team/pytorch_geometric/issues/9112

XinweiHe: > Can we test for [HGTConv](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.conv.HGTConv.html#torch_geometric.nn.conv.HGTConv) and [HanConv](https://pytorch-geometric.readthedocs.io/en/latest/generated/torch_geometric.nn.conv.HANConv.html#torch_geometric.nn.conv.HANConv), to ensure that it works for Hetero GNNs not created by the to_hetero operation.

`HGTConv` and `HanConv` are not relevant here, because `AttentionExplainer` only works with GNN model with attentions, in particular, `GATConv`, `GATv2Conv` or `TransformerConv`.

We will support `HeteroConv` with these attention GNN models in this PR though.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/10169?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) Report
Attention: Patch coverage is `89.13043%` with `10 lines` in your changes missing coverage. Please review.
> Please [upload](https://docs.codecov.com/docs/codecov-uploader) report for BASE (`xinwei_support_pg_explainer_hetero_v3@744f830`). [Learn more](https://docs.codecov.io/docs/error-reference?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#section-missing-base-commit) about missing BASE report.

| [Files with missing lines](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/10169?dropdown=coverage&src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) | Patch % | Lines |
|---|---|---|
| [...geometric/explain/algorithm/attention\_explainer.py](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/10169?src=pr&el=tree&filepath=torch_geometric%2Fexplain%2Falgorithm%2Fattention_explainer.py&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team#diff-dG9yY2hfZ2VvbWV0cmljL2V4cGxhaW4vYWxnb3JpdGhtL2F0dGVudGlvbl9leHBsYWluZXIucHk=) | 89.13% | [10 Missing :warning: ](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/10169?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team) |

<details><summary>Additional details and impacted files</summary>


```diff
@@                           Coverage Diff                            @@
##             xinwei_support_pg_explainer_hetero_v3   #10169   +/-   ##
========================================================================
  Coverage                                         ?   85.36%           
========================================================================
  Files                                            ?      496           
  Lines                                            ?    33840           
  Branches                                         ?        0           
========================================================================
  Hits                                             ?    28889           
  Misses                                           ?     4951           
  Partials                                         ?        0           
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/pyg-team/pytorch_geometric/pull/10169?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=pyg-team).

<details><summary> :rocket: New features to boost your workflow: </summary>

- :snowflake: [Test Analytics](https://docs.codecov.com/docs/test-analytics): Detect flaky tests, report on failures, and find test suite problems.
</details>
wsad1: > HGTConv and HanConv are not relevant here, because AttentionExplainer only works with GNN model with attentions, in particular, GATConv, GATv2Conv or TransformerConv.

`HanConv` and `HGTConv` use attention while message passing.
XinweiHe: > HanConv and HGTConv use attention while message passing.

You are right but at the moment, attention explainer is only designed to handle the case where `alpha` exists as an input parameter in `def message(...)` function. Since this is not the setup for `HanConv` and `HGTConv` at the moment we would need to spend extra effort supporting them which would probably be the best to not mix with this initial heterogeneous graph support for attention explainer.
XinweiHe: could we get this effort merge and create a follow-up ticket to handle the `HanConv` and `HGTConv` cases?

Patch:
diff --git a/CHANGELOG.md b/CHANGELOG.md
index 6be207b19..7fc4614d6 100644
--- a/CHANGELOG.md
+++ b/CHANGELOG.md
@@ -8,6 +8,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 ### Added
 
 - Added PyTorch 2.6 support ([#10170](https://github.com/pyg-team/pytorch_geometric/pull/10170))
+- Added support for heterogenous graphs in `AttentionExplainer` ([#10169](https://github.com/pyg-team/pytorch_geometric/pull/10169))
 - Added support for heterogenous graphs in `PGExplainer` ([#10168](https://github.com/pyg-team/pytorch_geometric/pull/10168))
 - Added support for heterogenous graphs in `GNNExplainer` ([#10158](https://github.com/pyg-team/pytorch_geometric/pull/10158))
 - Added Graph Positional and Structural Encoder (GPSE) and example ([#9018](https://github.com/pyg-team/pytorch_geometric/pull/9018)) ([#10118](https://github.com/pyg-team/pytorch_geometric/pull/10118))
diff --git a/torch_geometric/explain/algorithm/attention_explainer.py b/torch_geometric/explain/algorithm/attention_explainer.py
index 69f2ebeb1..af7cae47d 100644
--- a/torch_geometric/explain/algorithm/attention_explainer.py
+++ b/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:


Test patch:
diff --git a/test/explain/algorithm/test_attention_explainer.py b/test/explain/algorithm/test_attention_explainer.py
index e4479367d..6167dcdff 100644
--- a/test/explain/algorithm/test_attention_explainer.py
+++ b/test/explain/algorithm/test_attention_explainer.py
@@ -1,9 +1,25 @@
 import pytest
 import torch
 
-from torch_geometric.explain import AttentionExplainer, Explainer
-from torch_geometric.explain.config import ExplanationType, MaskType
-from torch_geometric.nn import AttentiveFP, GATConv, GATv2Conv, TransformerConv
+from torch_geometric.explain import (
+    AttentionExplainer,
+    Explainer,
+    HeteroExplanation,
+)
+from torch_geometric.explain.config import (
+    ExplanationType,
+    MaskType,
+    ModelConfig,
+    ModelMode,
+)
+from torch_geometric.nn import (
+    AttentiveFP,
+    GATConv,
+    GATv2Conv,
+    TransformerConv,
+    to_hetero,
+)
+from torch_geometric.nn.conv import HeteroConv
 
 
 class AttentionGNN(torch.nn.Module):
@@ -20,6 +36,141 @@ class AttentionGNN(torch.nn.Module):
         return x
 
 
+class HeteroAttentionGNN(torch.nn.Module):
+    def __init__(self, metadata, model_config=None):
+        super().__init__()
+        self.model_config = model_config
+
+        # Create a single BaseGNN that uses all three attention mechanisms
+        class BaseGNN(torch.nn.Module):
+            def __init__(self):
+                super().__init__()
+                # Use different attention mechanisms in sequence
+                self.conv1 = GATConv((-1, -1), 16, heads=2,
+                                     add_self_loops=False)
+                self.conv2 = GATv2Conv((-1, -1), 16, heads=2,
+                                       add_self_loops=False)
+                self.conv3 = TransformerConv((-1, -1), 32, heads=1)
+
+            def forward(self, x, edge_index):
+                x = self.conv1(x, edge_index).relu()
+                x = self.conv2(x, edge_index).relu()
+                x = self.conv3(x, edge_index)
+                return x
+
+        # Convert to heterogeneous model with a single to_hetero call
+        self.gnn = to_hetero(BaseGNN(), metadata, debug=False)
+
+        # Output dimension based on model config
+        out_channels = 7 if (model_config and model_config.mode
+                             == ModelMode.multiclass_classification) else 1
+        self.lin = torch.nn.Linear(32, out_channels)
+
+    def forward(self, x_dict, edge_index_dict, **kwargs):
+        # Process through the heterogeneous GNN
+        out_dict = self.gnn(x_dict, edge_index_dict)
+
+        # Project paper node embeddings for classification/regression
+        x = self.lin(out_dict['paper'])
+
+        # Apply appropriate output transformation based on model config
+        if self.model_config:
+            if self.model_config.mode == ModelMode.binary_classification:
+                if self.model_config.return_type == 'probs':
+                    x = x.sigmoid()
+            elif self.model_config.mode == ModelMode.multiclass_classification:
+                if self.model_config.return_type == 'probs':
+                    x = x.softmax(dim=-1)
+                elif self.model_config.return_type == 'log_probs':
+                    x = x.log_softmax(dim=-1)
+
+        return x
+
+
+class HeteroConvAttentionGNN(torch.nn.Module):
+    def __init__(self, metadata, model_config=None):
+        super().__init__()
+        self.model_config = model_config
+
+        # Determine output channels based on model_config
+        self.out_channels = 1
+        if (model_config
+                and model_config.mode == ModelMode.multiclass_classification):
+            self.out_channels = 7
+
+        # Initialize node type-specific layers
+        self.lin_dict = torch.nn.ModuleDict()
+        self.initialized = False
+
+        # Create a dictionary of attention-based convolutions for each edge
+        # type
+        conv_dict = {}
+        for edge_type in metadata[1]:  # metadata[1] contains edge types
+            src_type, _, dst_type = edge_type
+            if src_type == dst_type:
+                # For same node type, use GATConv with add_self_loops=False
+                # Use concat=False to avoid dimension issues
+                conv_dict[edge_type] = GATConv(
+                    (-1, -1), 32, heads=2, add_self_loops=False, concat=False)
+            else:
+                # For different node types, use GATv2Conv with
+                # add_self_loops=False Use concat=False to avoid dimension
+                # issues
+                conv_dict[edge_type] = GATv2Conv(
+                    (-1, -1), 32, heads=2, add_self_loops=False, concat=False)
+
+        # Create the HeteroConv layer
+        self.conv = HeteroConv(conv_dict, aggr='sum')
+
+        # Output layer will be initialized in forward pass
+        self.out_lin = None
+
+    def _initialize_layers(self, x_dict):
+        """Initialize layers with correct dimensions when we first see the
+        data.
+        """
+        if not self.initialized:
+            # Initialize input projections
+            for node_type, x in x_dict.items():
+                in_channels = x.size(-1)
+                self.lin_dict[node_type] = torch.nn.Linear(in_channels,
+                                                           32).to(x.device)
+
+            # Initialize output projection
+            self.out_lin = torch.nn.Linear(32, self.out_channels).to(
+                x_dict['paper'].device)
+
+            self.initialized = True
+
+    def forward(self, x_dict, edge_index_dict):
+        # Initialize layers if not done yet
+        self._initialize_layers(x_dict)
+
+        # Apply node type-specific transformations
+        h_dict = {}
+        for node_type, x in x_dict.items():
+            h_dict[node_type] = self.lin_dict[node_type](x).relu_()
+
+        # Apply heterogeneous convolution
+        out_dict = self.conv(h_dict, edge_index_dict)
+
+        # Final transformation for paper nodes
+        out = self.out_lin(out_dict['paper'])
+
+        # Apply transformations based on model_config if available
+        if self.model_config:
+            if self.model_config.mode == ModelMode.binary_classification:
+                if self.model_config.return_type == 'probs':
+                    out = out.sigmoid()
+            elif self.model_config.mode == ModelMode.multiclass_classification:
+                if self.model_config.return_type == 'probs':
+                    out = out.softmax(dim=-1)
+                elif self.model_config.return_type == 'log_probs':
+                    out = out.log_softmax(dim=-1)
+
+        return out
+
+
 x = torch.randn(8, 3)
 edge_index = torch.tensor([
     [0, 1, 1, 2, 2, 3, 3, 4, 4, 5, 5, 6, 6, 7],
@@ -82,3 +233,80 @@ def test_attention_explainer_attentive_fp(check_explanation):
 
     explanation = explainer(x, edge_index, edge_attr=edge_attr, batch=batch)
     check_explanation(explanation, None, explainer.edge_mask_type)
+
+
+@pytest.mark.parametrize('index', [None, 2, torch.arange(3)])
+def test_attention_explainer_hetero(index, hetero_data,
+                                    check_explanation_hetero):
+    # Create model configuration
+    model_config = ModelConfig(
+        mode='multiclass_classification',
+        task_level='node',
+        return_type='raw',
+    )
+
+    # Get metadata from hetero_data
+    metadata = hetero_data.metadata()
+
+    # Create the hetero attention model
+    model = HeteroAttentionGNN(metadata, model_config)
+
+    # Create the explainer
+    explainer = Explainer(
+        model=model,
+        algorithm=AttentionExplainer(),
+        explanation_type='model',
+        edge_mask_type='object',
+        model_config=model_config,
+    )
+
+    # Generate the explanation
+    explanation = explainer(
+        hetero_data.x_dict,
+        hetero_data.edge_index_dict,
+        index=index,
+    )
+
+    # Check that the explanation is correct
+    assert isinstance(explanation, HeteroExplanation)
+    check_explanation_hetero(explanation, None, explainer.edge_mask_type,
+                             hetero_data)
+
+
+@pytest.mark.parametrize('index', [None, 2, torch.arange(3)])
+def test_attention_explainer_hetero_conv(index, hetero_data,
+                                         check_explanation_hetero):
+    """Test AttentionExplainer with HeteroConv using attention-based layers."""
+    # Create model configuration
+    model_config = ModelConfig(
+        mode='multiclass_classification',
+        task_level='node',
+        return_type='raw',
+    )
+
+    # Get metadata from hetero_data
+    metadata = hetero_data.metadata()
+
+    # Create the hetero conv attention model
+    model = HeteroConvAttentionGNN(metadata, model_config)
+
+    # Create the explainer
+    explainer = Explainer(
+        model=model,
+        algorithm=AttentionExplainer(),
+        explanation_type='model',
+        edge_mask_type='object',
+        model_config=model_config,
+    )
+
+    # Generate the explanation
+    explanation = explainer(
+        hetero_data.x_dict,
+        hetero_data.edge_index_dict,
+        index=index,
+    )
+
+    # Check that the explanation is correct
+    assert isinstance(explanation, HeteroExplanation)
+    check_explanation_hetero(explanation, None, explainer.edge_mask_type,
+                             hetero_data)