Fix fidelity metric and assigning edge masks as tensors

CHANGELOG.md

@@ -38,7 +38,7 @@ The format is based on [Keep a Changelog](http://keepachangelog.com/en/1.0.0/).
 - Added the `AirfRANS` dataset ([#6287](https://github.com/pyg-team/pytorch_geometric/pull/6287))
 - Added `AttentionExplainer` ([#6279](https://github.com/pyg-team/pytorch_geometric/pull/6279))
 - Added (un)faithfulness explainability metric ([#6090](https://github.com/pyg-team/pytorch_geometric/pull/6090))
-- Added fidelity explainability metric ([#6116](https://github.com/pyg-team/pytorch_geometric/pull/6116))
+- Added fidelity explainability metric ([#6116](https://github.com/pyg-team/pytorch_geometric/pull/6116), [#6510](https://github.com/pyg-team/pytorch_geometric/pull/6510))
 - Added subgraph visualization of GNN explanations ([#6235](https://github.com/pyg-team/pytorch_geometric/pull/6235), [#6271](https://github.com/pyg-team/pytorch_geometric/pull/6271))
 - Added weighted negative sampling option in `LinkNeighborLoader` ([#6264](https://github.com/pyg-team/pytorch_geometric/pull/6264))
 - Added the `BA2MotifDataset` explainer dataset ([#6257](https://github.com/pyg-team/pytorch_geometric/pull/6257))

torch_geometric/explain/algorithm/utils.py

@@ -1,15 +1,16 @@
-from typing import Dict
+from typing import Dict, Union
 
 import torch
 from torch import Tensor
+from torch.nn import Parameter
 
 from torch_geometric.nn import MessagePassing
 from torch_geometric.typing import EdgeType
 
 
 def set_masks(
     model: torch.nn.Module,
-    mask: Tensor,
+    mask: Union[Tensor, Parameter],
     edge_index: Tensor,
     apply_sigmoid: bool = True,
 ):
@@ -19,6 +20,14 @@ def set_masks(
     # Loop over layers and set masks on MessagePassing layers:
     for module in model.modules():
         if isinstance(module, MessagePassing):
+
+            # Convert mask to a param if it was previously registered as one.
+            # This is a workaround for the fact that PyTorch does not allow
+            # assignments of pure tensors to parameter attributes:
+            if (not isinstance(mask, Parameter)
+                    and '_edge_mask' in module._parameters):
+                mask = Parameter(mask)
+
             module.explain = True
             module._edge_mask = mask
             module._loop_mask = loop_mask
@@ -27,7 +36,7 @@ def set_masks(
 
 def set_hetero_masks(
     model: torch.nn.Module,
-    mask_dict: Dict[EdgeType, Tensor],
+    mask_dict: Dict[EdgeType, Union[Tensor, Parameter]],
     edge_index_dict: Dict[EdgeType, Tensor],
     apply_sigmoid: bool = True,
 ):

torch_geometric/explain/metric/fidelity.py

@@ -20,7 +20,7 @@ def fidelity(
     Fidelity evaluates the contribution of the produced explanatory subgraph
     to the initial prediction, either by giving only the subgraph to the model
     (fidelity-) or by removing it from the entire graph (fidelity+).
-    The fidelity scores capture how good an explanable model reproduces the
+    The fidelity scores capture how good an explainable model reproduces the
     natural phenomenon or the GNN model logic.
 
     For **phenomenon** explanations, the fidelity scores are given by: