(self, x_j, edge_attr)
| 61 | return out |
| 62 | |
| 63 | def apply_edge(self, x_j, edge_attr): |
| 64 | edge_attr, idx = tf.unique(edge_attr) |
| 65 | matrix = tf.gather(self.matrix, edge_attr) |
| 66 | matrix = tf.gather(matrix, idx) |
| 67 | matrix = tf.reshape(matrix, [-1, self.dim, self.fea_dim]) |
| 68 | x_j = tf.expand_dims(x_j, -1) |
| 69 | res = tf.matmul(matrix, x_j) |
| 70 | return tf.reshape(res, [-1, self.dim]) |
| 71 | |
| 72 | def apply_node(self, x, aggr_out): |
| 73 | return self.fc(x) + aggr_out |