Confused about how dimensionality works in Heterogeneous Graph Learning

[enter802]

I was looking through the implementation examples of Heterogeneous graph learning here, and I implemented the model successfully for my problem. However, I am unable to explain to myself how the dimensionality works. So for example, the network below

from torch_geometric.nn import GATConv, Linear, to_hetero
class GAT(torch.nn.Module):
    def __init__(self, hidden_channels, out_channels):
        super().__init__()
        self.conv1 = GATConv((-1, -1), hidden_channels, add_self_loops=False)
        self.lin1 = Linear(-1, hidden_channels)
        self.conv2 = GATConv((-1, -1), out_channels, add_self_loops=False)
        self.lin2 = Linear(-1, out_channels)

    def forward(self, x, edge_index):
        x = self.conv1(x, edge_index) + self.lin1(x)
        x = x.relu()
        x = self.conv2(x, edge_index) + self.lin2(x)
        return x
model = GAT(hidden_channels=64, out_channels=dataset.num_classes)
model = to_hetero(model, data.metadata(), aggr='sum')

in the forward function, in the line

x = self.conv1(x, edge_index) + self.lin1(x)

I am not sure how the Linear module takes in the features as input and then added to the GATConv module. Is it that for the Linear layer, the features are flattened and passed through it?

I am not sure if I am clear in explaining my question. Hope to get some feedback soon. Thanks!

[wsad1]

In the example you shared, the input dimension of self.lin1 and self.lin2 are initialized by to_hetero, to the in_channels of the nodes. So the features don't need to be flattened while passing to self.lin1 or self.lin2.

[enter802]

Do you think this would be a good representation of the lines

x = self.conv1(x, edge_index) + self.lin1(x)
x = x.relu()

?

[wsad1]

Yes that looks good.