| 124 | return score |
| 125 | |
| 126 | def multi_head(self, query, key, value): |
| 127 | query = self.lin_q(query) |
| 128 | key = self.lin_k(key) |
| 129 | value = self.lin_v(value) |
| 130 | out_channels_per_head = self.dim // self.heads |
| 131 | |
| 132 | query_shape = [-1, tf.shape(query)[1], |
| 133 | self.heads, out_channels_per_head] |
| 134 | query = tf.transpose(tf.reshape(query, query_shape), [1, 0, 2, 3]) |
| 135 | |
| 136 | key_shape = [-1, tf.shape(key)[1], self.heads, out_channels_per_head] |
| 137 | key = tf.transpose(tf.reshape(key, key_shape), [1, 0, 2, 3]) |
| 138 | |
| 139 | value_shape = [-1, tf.shape(value)[1], |
| 140 | self.heads, out_channels_per_head] |
| 141 | value = tf.transpose(tf.reshape(value, value_shape), [1, 0, 2, 3]) |
| 142 | |
| 143 | out = self.attention(query, key, value) |
| 144 | out = tf.transpose(out, [1, 0, 2, 3]) |
| 145 | out_shape = [-1, tf.shape(query)[1], self.dim] |
| 146 | out = tf.reshape(out, out_shape) |
| 147 | return out |
| 148 | |
| 149 | def __call__(self, x, edge_index, size=None, **kwargs): |
| 150 | if isinstance(x, tf.Tensor): |