def _distance_logits(self, x, keys):
init = get_keras_initialization(self.init)
key_w = tf.get_variable("key_w", shape=keys.shape.as_list()[-1], initializer=init, dtype=tf.float32)
key_logits = tf.tensordot(keys, key_w, axes=[[2], [0]]) # (batch, key_len)
x_w = tf.get_variable("input_w", shape=x.shape.as_list()[-1], initializer=init, dtype=tf.float32)
x_logits = tf.tensordot(x, x_w, axes=[[2], [0]]) # (batch, x_len)
dot_w = tf.get_variable("dot_w", shape=x.shape.as_list()[-1], initializer=init, dtype=tf.float32)
# Compute x * dot_weights first, the batch mult with x
x_dots = x * tf.expand_dims(tf.expand_dims(dot_w, 0), 0)
dot_logits = tf.matmul(x_dots, keys, transpose_b=True)
return dot_logits + tf.expand_dims(key_logits, 1) + tf.expand_dims(x_logits, 2)
评论列表
文章目录