# How to implement edge feature in message passing algorithm of Pytorch Geometric?

I am trying to implement a custom message passing function $$x^k_i = x^{k-1}_i + m^k_i$$ where $$x_i$$ is the node embedding of the i'th node and $$m^k_i$$ is the message getting delivered at the i'th node from its neighbours.

The message is defined as $$m^k_i = \frac{1}{N} \sum_{j\in N(i)} ReLU(z_{ij})$$ , where $$z_{ij} = x^{k-1}i + x^{k-1}j + \sigma(w*e_{ij} + b)$$ ?

Where $$x_j$$ = neighbouring node and $$e_{ij} =$$ Edge feature between node 'i' and 'j'.

Code:

import torch
from torch.nn import Sequential as Seq, Linear, ReLU
from torch_geometric.nn import MessagePassing

class GC_a(MessagePassing):
def _init_(self, in_channels, out_channels, hidden_channel, edge_feature):
super(GC_a, self)._init_(aggr='mean')
self.lin_e = torch.nn.Linear(edge_feature, hidden_channel)
self.act_e = torch.nn.ReLU()

def forward(self, x, edge_index, edge_feature):
# x has shape [N, in_channels]
# edge_index has shape [2, E]
edge_index, _ = remove_self_loops(edge_index)

return self.propagate(edge_index, size=(x.size(0), x.size(0)), x=x)

def message(self, x_i, x_j, e_ij): # e_ij is edge feature ()

z_ij = torch.cat([x_i, x_j], dim=1) # These are the first 2 terms of the message expression; z_ij has shape [in_channels+in_channels,1]

e_f = self.lin_e(e_ij) # e_ij [edge_feature,1]
e_f = self.act_e(e_f) # [hidden_channel,1]

z_ij = torch.cat([z_ij, e_f], dim=1) # z_ij has shape [in_channels+in_channels+hidden_channels,1]

return z_ij

def update(self, aggr_out, x_i):

# aggr_out has shape [N, out_channels]
new_embedding = torch.cat([aggr_out], dim=1)

return new_embedding