gin conv for hetero data

[DogitoErgoSum]

Is there a way to apply gin Conv for hetero data like this?
For example, let 1,2,3 be 3 different node types with feature vectors of size A,B,C.

Apply 2 linear transformations(1 for each edge type) to A such that A becomes of the same shape as B and C to then be able to use GINConv. All in one layer

[rusty1s]

The GINConv does support bipartite graphs, and so it should also be applicable for HeteroData. Importantly, it does only work though for equal feature dimensionality across all types. You can ensure that by mapping all node types to the same feature space (e.g., by using distinct Linear layers before-hand).

[DogitoErgoSum]

Thank you!

[DogitoErgoSum]

Is there a simple way to do that?

[rusty1s]

from torch_geometric.nn import Linear

self.lins = torch.nn.ModuleDict()
for node_type in node_types:
    self.lins[node_type] = Linear(-1, hidden_channels)

def forward(self, x_dict, ...):
    x_dict = {self.lins[node_type](x) for node_type, x in x_dict.items()}