How does to_hetero work?

[shahinghasemi]

I’m confused about the theory behind to_hetero API. As written on the specs:

Converts a homogeneous GNN model into its heterogeneous equivalent in which node representations are learned for each node type in metadata[0], and messages are exchanged between each edge type in metadata[1], as denoted in the “Modeling Relational Data with Graph Convolutional Networks” paper

Let’s say that we have a simple model as below:

class GNN(torch.nn.Module):
    def __init__(self):
        self.conv1 = SAGEConv((-1, -1), 32)
    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index)
        return x
model = GNN()
model = to_hetero(model, metadata)

Now according to doc SAGEConv compute the representations as below:
SAGEConv
So is it correct to say that the model above (which combines SAGEConv and to_hetero) computes the representations as below:
Final Model

[rusty1s]

Yes that is correct. conv1 will be duplicated and applied for each edge type, and the result will be summed up.