GNNExplainer for HeteroData

[Lsabetta]

Hi!

I'm working with a custom GNN model that takes a HeteroData object as input, with two relations. I want to use GNNExplainer to understand the model's predictions. I successfully defined the explainer as follows:

from torch_geometric.explain import Explainer, GNNExplainer
from torch_geometric.explain.config import ModelConfig
explainer = Explainer(
    model=model,  # GNN model
    algorithm=GNNExplainer(epochs=200),  # Algorithm for explainability
    explanation_type="model",  # Explain the model prediction
    model_config=ModelConfig(
        mode="multiclass_classification",  # Model mode
        task_level="node",  # Node-level task
        return_type="raw"  # Model returns raw logits
    ),
    edge_mask_type="object"  
)

However, when I call it with a HeteroData object like this:

explanation = explainer(
    x=sub_data.x_dict,
    edge_index=sub_data.edge_index_dict,
    edge_attr=sub_data.edge_attr_dict,
    index=node_idx_to_explain
)

I get the following error:

ValueError: Heterogeneous graphs not yet supported in 'GNNExplainer'

Based on the PyTorch Geometric documentation, I expected Explainer to handle both homogeneous and heterogeneous graphs. I read on another discussion that support for HeteroData was foreseen for pyg 2,2.
Is this behavior expected for GNNExplainer, or am I missing something in my setup? If GNNExplainer doesn't support heterogeneous graphs, could you please recommend an alternative explainability method compatible with HeteroData?

Thanks in advance for your kind help!
Luigi