| 66 | return e_id |
| 67 | |
| 68 | def get_edge_attr(self, block): |
| 69 | n_id_dst = tf.cast(tf.expand_dims(block.n_id, -1), |
| 70 | dtype=tf.float32) |
| 71 | n_id_src= mp_ops.gather(n_id_dst, block.res_n_id) |
| 72 | n_id_src = mp_ops.gather(n_id_src, |
| 73 | block.edge_index[0]) |
| 74 | n_id_dst = mp_ops.gather(n_id_dst, |
| 75 | block.edge_index[1]) |
| 76 | n_id_src = tf.cast(tf.squeeze(n_id_src, -1), dtype=tf.int64) |
| 77 | n_id_dst = tf.cast(tf.squeeze(n_id_dst, -1), dtype=tf.int64) |
| 78 | edge_attr = self.to_edge(n_id_src, n_id_dst, block.e_id) |
| 79 | return edge_attr |
| 80 | |
| 81 | def calculate_conv(self, conv, inputs, edge_index, |
| 82 | size=None, edge_attr=None): |