Question about using to_hetero and HeteroConv with a MessagePassing module that updates nodes and edges

[thebookins]

I have a MessagePassing module (we'll call it EdgeUpdatingConv) that on its forward pass updates nodes and edges, returning the tuple x, edge_attr.

Say I want to use this layer num_layers times, with the output of each layer fed into the next layer as input. Furthermore, I want to apply these layers to a heterogenous graph. According to https://pytorch-geometric.readthedocs.io/en/2.6.0/notes/heterogeneous.html this can be accomplished by applying the to_hetero transformation to the module, or by using HeteroConv. However, both to_hetero and HeteroConv assume that the MessagePassing convolutions return a single tensor representing the updated nodes.

For example, the following code will fail at the aggregation step:

import torch
from torch_geometric.nn import HeteroConv

class MyHeteroGraphNet(torch.nn.Module):
    def __init__(self, hidden_dim, num_layers, metadata):
        super(MyHeteroGraphNet, self).__init__()

        self.convs = torch.nn.ModuleList()
        for _ in range(num_layers):
            conv = HeteroConv({
                edge_type: EdgeUpdatingConv(-1, hidden_dim)
                for edge_type in metadata[1]
            }, aggr='sum')
            self.convs.append(conv)

    def forward(self, x_dict, edge_index_dict, edge_attr_dict):
        for conv in self.convs:
            x_dict, edge_attr_dict = conv(x_dict, edge_index_dict, edge_attr_dict)
        return x_dict

I have managed to get the right result by writing a custom HeteroConv class that expects each convolution to return updated nodes and edges (and the custom class itself returns a dictionary of updated nodes and a dictionary of updated edges).

However, I'm curious if there is a better way of achieving this, or if it makes sense in the future for to_hetero and HeteroConv to support edge-updating convolutions.