heterogeneous graph with message using edge_attr

[choiwanuk]

Hello! I'm trying to creating a heterogeneous graph NN with user-defined message like under codes.

is it possible to coding heterogeneous graph NN with edge_attr like below?

class InteractionNetwork(MessagePassing):
    def __init__(self):
        super(InteractionNetwork, self).__init__(aggr='add')
        self.node_fn = LazyLinear(128*2,128)
        self.edge_fn = LazyLinear(128*3,128)
        
    def reset_parameters(self):
        self.node_fn.reset_parameters()
        self.edge_fn.reset_parameters()
        
    def forward(self, x, edge_index, edge_attr):
        x_residual = x
        edge_attr_residual = edge_attr
        x, edge_attr = self.propagate(edge_index=edge_index, x=x, edge_attr=edge_attr)
        return x+x_residual, edge_attr+edge_attr_residual

    def message(self, x_i, x_j, edge_attr):
        edge_attr = torch.cat([x_i, x_j, edge_attr], dim=-1)
        edge_attr = self.edge_fn(edge_attr)
        return edge_attr

    def update(self, x_updated, x, edge_attr):
        x_updated = torch.cat([x_updated, x], dim=-1)
        x_updated = self.node_fn(x_updated)
        return x_updated, edge_attr
    

class graphNet(torch.nn.Module):
    def __init__(self, in_channels,out_channels):
        super().__init__()
        self.node_embedding = nn.Linear(out_channels,in_channels)
        self.edge_embedding = nn.Linear(out_channels,in_channels)
        self.conv1 = InteractionNetwork()
        self.lin = nn.Linear(in_channels, out_channels, bias=True)

    def forward(self, x, edge_index, edge_attr):
        embedded_x = self.node_embedding(x)
        embedded_edge_attr = self.edge_embedding(edge_attr)
        output_x, output_edge_attr = self.conv1(embedded_x, edge_index, embedded_edge_attr)
        return self.lin(output_x)
        #return embedded_x

I converted it with to_hetero function and put a dummy hetero graphs but error was shown.

hope someone to tell me advice.. please.

I referenced test code for Hetero to design..

pytorch_geometric/test/nn/test_to_hetero_transformer.py

Line 24 in df1d3f5

class Net1(torch.nn.Module):

[rusty1s]

The issue is that in update, x is a tuple of (source_node_features, dst_node_features) in the heterogeneous graph case. You could do something like

x = x[1] if isinstance(x, tuple) else x

or perform the update step just after self.propagate.