Can the message function take the node and all its neighbors as input?

[gebrahimi91]

Hi,

I am trying to use my own designed network to message and aggregate the messages. The way that i aggregate the inputs is not simple mean, sum, etc. so node i representation and all its neighbors are fed to a neural network together and the aggregation happens in the neural network and not after message from each neighbor is created. in other words, i need a message function like this:

def message(x_i):
return neural network([x_i, x_j for all neighbors of x_i])
after doing this the message is already passed and aggregated (i do not need any more aggregation and aggregation can be set to None).
Is there an easy way to do it using message function or message_and_update function?

Thanks.

[HongtaoYang]

I think you need to

1. implement the message_and_aggregate method and call your neural_network module from there. This way you control both "message" and "aggregate" process at the same time. This should help you because in your case these two processes are not easily separable.
2. make sure the edge_index is a SparseTensor so that message_and_aggregate method gets executed and not message method.

[rusty1s]

Currently, this is only possible via some hacks:

def message_and_aggregate(self, adj_t, x):
    col = adj_t.storage.col()
    rowcount = adj_t.storage.rowcount()
    x_j = x[col].split(rowcount.tolist())
    x_neighbor = torch.nn.utils.rnn.pad_sequence(x_j, batch_first=True)
    # x_neighbor will have shape [num_nodes, num_max_neighbors, num_features]
    ...  # Do whatever you like to do with this