LinkX implementation

[Chengwei94]

Hi there, I was looking at the linkx implementation in pytorch geometric code, and the formula, and find that the implementation seems to be different from the formula. Is there a reason for it ?

def forward(self, x: OptTensor, edge_index: Adj,
            edge_weight: OptTensor = None) -> Tensor:
    """"""
    out = self.edge_lin(edge_index, edge_weight)
    if self.num_edge_layers > 1:
        out = out.relu_()
        out = self.edge_norm(out)
        out = self.edge_mlp(out)

    out = out + self.cat_lin1(out)

    if x is not None:
        x = self.node_mlp(x)
        out = out + x
        out = out + self.cat_lin2(x)

    return self.final_mlp(out.relu_())

vs

[Chengwei94]

After giving it more thought, I realise that they are equivalent. Y = MLP(\sigma(W_A[H_A] + W_X[H_X] + H_A + H_X))