to_hetero throws NotImplementedError for certain metadata

[wgeul]

Hi,

I'm playing around with a heterogeneous graph structure that apparently fails when calling to_hetero. I'm not able to derive what the issue is.
Do you know why (2) this doesn't work?

from torch.nn import ModuleList
import torch
from torch_geometric.nn import SAGEConv, to_hetero

class GNNEncoder(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels, nlayers):
        super().__init__()
        # Block
        self.convblock = ModuleList(
            [SAGEConv((-1, -1), hidden_channels) for _ in range(nlayers)]
        )
        self.conv2 = SAGEConv((-1, -1), out_channels)

    def forward(self, x, edge_index):
        for conv in self.convblock:
            x = conv(x, edge_index).relu()
        x = self.conv2(x, edge_index)
        return x

(1) works fine:

metadata = [["a", "b"], [("a", "to", "b"), ("b", "rev_to", "a")]]
encoder = GNNEncoder(32, 32, 1)
encoder = to_hetero(encoder, metadata, aggr="sum")

(2) throws an error:

metadata = (
    [
        "a",
        "b",
        "c",
        "d",
        "e",
        "f",
        "g",
        "h",
    ],
    [
        ("a", "to", "b"),
        ("b", "to", "f"),
        ("f", "to", "b"),
        ("a", "to", "e"),
        ("b", "to", "c"),
        ("b", "to", "d"),
        ("c", "to", "h"),
        ("b", "rev_to", "a"),
        ("f", "rev_to", "b"),
        ("b", "rev_to", "f"),
        ("e", "rev_to", "a"),
        ("c", "rev_to", "b"),
        ("d", "rev_to", "b"),
        ("h", "rev_to", "c"),
    ],
)

encoder = GNNEncoder(32, 32, 1)
encoder = to_hetero(encoder, metadata, aggr="sum")

error:

NotImplementedError                       Traceback (most recent call last)
[test_model_sageconv.py) in ()
     [54](test_model_sageconv.py?line=53) 
     [55](test_model_sageconv.py?line=54) encoder = GNNEncoder(32, 32, 1)
---> [56](test_model_sageconv.py?line=55) encoder = to_hetero(encoder, metadata, aggr="sum")

[..python3.10/site-packages/torch_geometric/nn/to_hetero_transformer.py) in to_hetero(module, metadata, aggr, input_map, debug)
    116     """
    117     transformer = ToHeteroTransformer(module, metadata, aggr, input_map, debug)
--> 118     return transformer.transform()
    119 
    120 

[..python3.10/site-packages/torch_geometric/nn/fx.py) in transform(self)
    155             elif is_global_pooling_op(self.module, op, node.target):
    156                 op = 'call_global_pooling_module'
--> 157             getattr(self, op)(node, node.target, node.name)
    158 
    159         # Remove all unused nodes in the computation graph, i.e., all nodes

[..python3.10/site-packages/torch_geometric/nn/to_hetero_transformer.py) in call_method(self, node, target, name)
    292         self.graph.inserting_after(node)
    293         for key in self.metadata[int(self.is_edge_level(node))]:
--> 294             args, kwargs = self.map_args_kwargs(node, key)
...
--> 387                     raise NotImplementedError
    388             elif isinstance(value, dict):
    389                 return {k: _recurse(v) for k, v in value.items()

[wgeul]

Found it, g was not present in the edges. Apologies.