def testDotProductUnMaskedAttentionRelativeV2(self):
x = np.random.rand(5, 7, 12, 32)
y = np.random.rand(5, 7, 12, 32)
a = common_attention.dot_product_unmasked_self_attention_relative_v2(
tf.constant(x, dtype=tf.float32), tf.constant(y, dtype=tf.float32),
tf.constant(y, dtype=tf.float32), None, 35)
self.evaluate(tf.global_variables_initializer())
res = self.evaluate(a)
self.assertEqual(res.shape, (5, 7, 12, 32))
def testRelativeAttentionV2UnmaskedRelativeLargerThanLength(self):
# (batch, heads, length, depth)
x = np.random.rand(5, 4, 3, 7)
y = np.random.rand(5, 4, 3, 7)
max_relative_position = 16
a = common_attention.dot_product_unmasked_self_attention_relative_v2(
tf.constant(x, dtype=tf.float32),
tf.constant(y, dtype=tf.float32),
tf.constant(y, dtype=tf.float32),
None,
max_relative_position=max_relative_position,
heads_share_relative_embedding=False)
self.evaluate(tf.global_variables_initializer())
res = self.evaluate(a)
self.assertEqual(res.shape, (5, 4, 3, 7))
