| 270 | |
| 271 | |
| 272 | class EdgeEmbedding(MessagePassing): |
| 273 | |
| 274 | def __init__(self, num_rbf, hidden_channels): |
| 275 | super(EdgeEmbedding, self).__init__(aggr=None) |
| 276 | self.edge_proj = nn.Linear(num_rbf, hidden_channels) |
| 277 | |
| 278 | self.reset_parameters() |
| 279 | |
| 280 | def reset_parameters(self): |
| 281 | nn.init.xavier_uniform_(self.edge_proj.weight) |
| 282 | self.edge_proj.bias.data.fill_(0) |
| 283 | |
| 284 | def forward(self, edge_index, edge_attr, x): |
| 285 | # propagate_type: (x: Tensor, edge_attr: Tensor) |
| 286 | out = self.propagate(edge_index, x=x, edge_attr=edge_attr) |
| 287 | return out |
| 288 | |
| 289 | def message(self, x_i, x_j, edge_attr): |
| 290 | return (x_i + x_j) * self.edge_proj(edge_attr) |
| 291 | |
| 292 | def aggregate(self, features, index): |
| 293 | # no aggregate |
| 294 | return features |