Can nn.MessagePassing.message() pass message for multiple attributes?

[ipcamit]

I was trying to implement EGNN as in http://arxiv.org/abs/2102.09844.
The model is defined as (Sorry for poor formatting, github markdown was not rendering latex.)

𝐦𝑖𝑗=𝜙𝑒(𝐡𝑖,𝐡𝑗,‖𝑟𝑖−𝑟𝑗‖,𝑎𝑖𝑗)
𝑟𝑖=𝑟𝑖+𝐶∑(𝑟𝑖−𝑟𝑗)𝜙𝑟(𝐦𝑖𝑗)
𝐦𝑖=∑_(𝑖≠𝑗)𝐦𝑖𝑗
𝐡𝑖=𝜙ℎ(𝐡𝑖,𝐦𝑖)

Updating 𝐦𝑖𝑗 is trivial, I was wondering if message could also propagate, 𝑟𝑖 in single step? Currently if my message function returns more then 1 value I get error regarding dimensions of tuple.

Shifting the 𝑟𝑖 step to update means 1. 𝐡𝑖 is no longer lifted (its dimensions are same as that of number of nodes), and 2. I need to repeat distance and norm calculations.

[rusty1s]

In general, returning multiple outputs within message passing should be supported, see here for a simple test case. Note that aggregate(inputs, ...) will now hold a tuple, such that you may need to override aggregate as well.