code: from parameterized import parameterized
  
  import tensorflow as tf
  import numpy as np
  
  from opennmt.layers import transformer
  
  
  class TransformerTest(tf.test.TestCase):
  
    @parameterized.expand([[tf.bool], [tf.float32]])
    def testBuildFutureMask(self, dtype):
      length = [2, 4, 3]
      expected = np.array([
          [[1, 0, 0, 0],
           [1, 1, 0, 0],
           [1, 1, 0, 0],
           [1, 1, 0, 0]],
          [[1, 0, 0, 0],
           [1, 1, 0, 0],
           [1, 1, 1, 0],
           [1, 1, 1, 1]],
          [[1, 0, 0, 0],
           [1, 1, 0, 0],
           [1, 1, 1, 0],
           [1, 1, 1, 0]]]).astype(dtype.as_numpy_dtype)
  
      mask = transformer.future_mask(tf.constant(length), dtype=dtype)
      self.assertIs(mask.dtype, dtype)
      mask = self.evaluate(mask)
      self.assertTupleEqual(mask.shape, (len(length), max(length), max(length)))
      self.assertAllEqual(mask, expected)
  
    @parameterized.expand([[tf.bool], [tf.float32]])
    def testBuildFutureMaskWithMaxLen(self, dtype):
      length = [2, 4, 3]
      maximum_length = 5
      expected = np.array([
          [[1, 0, 0, 0, 0],
           [1, 1, 0, 0, 0],
           [1, 1, 0, 0, 0],
           [1, 1, 0, 0, 0],
           [1, 1, 0, 0, 0]],
          [[1, 0, 0, 0, 0],
           [1, 1, 0, 0, 0],
           [1, 1, 1, 0, 0],
           [1, 1, 1, 1, 0],
           [1, 1, 1, 1, 0]],
          [[1, 0, 0, 0, 0],
           [1, 1, 0, 0, 0],
           [1, 1, 1, 0, 0],
           [1, 1, 1, 0, 0],
           [1, 1, 1, 0, 0]]]).astype(dtype.as_numpy_dtype)
  
      mask = transformer.future_mask(
          tf.constant(length), maximum_length=maximum_length, dtype=dtype)
      self.assertIs(mask.dtype, dtype)
      mask = self.evaluate(mask)
      self.assertTupleEqual(mask.shape, (len(length), maximum_length, maximum_length))
      self.assertAllEqual(mask, expected)
  
    def testSplitHeads(self):
      batch_size = 3
      length = [5, 3, 7]
      num_heads = 8
      depth = 20
  
      inputs = tf.random.normal([batch_size, max(length), depth * num_heads], dtype=tf.float32)
      outputs = transformer.split_heads(inputs, num_heads)
  
      static_shape = outputs.shape
      self.assertEqual(num_heads, static_shape[1])
      self.assertEqual(depth, static_shape[-1])
      outputs = self.evaluate(outputs)
      self.assertAllEqual([batch_size, num_heads, max(length), depth], outputs.shape)
  
    def testCombineHeads(self):
      batch_size = 3
      length = [5, 3, 7]
      num_heads = 8
      depth = 20
  
      inputs = tf.random.normal([batch_size, num_heads, max(length), depth], dtype=tf.float32)
      outputs = transformer.combine_heads(inputs)
  
      static_shape = outputs.shape
      self.assertEqual(depth * num_heads, static_shape[-1])
      outputs = self.evaluate(outputs)
      self.assertAllEqual([batch_size, max(length), depth * num_heads], outputs.shape)
  
    def testSplitAndCombineHeads(self):
      batch_size = 3
      length = [5, 3, 7]
      num_heads = 8
      depth = 20
  
      inputs = tf.random.normal([batch_size, max(length), depth * num_heads], dtype=tf.float32)
      split = transformer.split_heads(inputs, num_heads)
      combined = transformer.combine_heads(split)
      inputs, combined = self.evaluate([inputs, combined])
      self.assertAllEqual(inputs, combined)
  
    def testRelativePositions(self):
      positions = transformer.relative_positions(4, 2)
      self.assertAllEqual(
          self.evaluate(positions),
          [[2, 3, 4, 4], [1, 2, 3, 4], [0, 1, 2, 3], [0, 0, 1, 2]])
  
    def testFeedForwardNetwork(self):
      ffn = transformer.FeedForwardNetwork(20, 10)
      x = tf.random.uniform([4, 5, 10])
      y = ffn(x)
      self.assertEqual(y.shape, x.shape)
  
    def testMultiHeadSelfAttention(self):
      attention = transformer.MultiHeadAttention(4, 20)
      queries = tf.random.uniform([4, 5, 10])
      mask = tf.sequence_mask([4, 3, 5, 2])
      context, _ = attention(queries, mask=mask)
      self.assertListEqual(context.shape.as_list(), [4, 5, 20])
  
    def testMultiHeadSelfAttentionWithCache(self):
      cache = (tf.zeros([4, 4, 0, 5]), tf.zeros([4, 4, 0, 5]))
      attention = transformer.MultiHeadAttention(4, 20)
      x = tf.random.uniform([4, 1, 10])
      _, cache = attention(x, cache=cache)
      self.assertEqual(cache[0].shape[2], 1)
      self.assertEqual(cache[1].shape[2], 1)
      _, cache = attention(x, cache=cache)
      self.assertEqual(cache[0].shape[2], 2)
      self.assertEqual(cache[1].shape[2], 2)
  
    def testMultiHeadSelfAttentionRelativePositions(self):
      attention = transformer.MultiHeadAttention(4, 20, maximum_relative_position=6)
      x = tf.random.uniform([2, 9, 10])
      mask = tf.sequence_mask([9, 7])
      y = attention(x, mask=mask)
  
    def testMultiHeadSelfAttentionRelativePositionsWithCache(self):
      attention = transformer.MultiHeadAttention(4, 20, maximum_relative_position=6)
      x = tf.random.uniform([4, 1, 10])
      cache = (tf.zeros([4, 4, 0, 5]), tf.zeros([4, 4, 0, 5]))
      _, cache = attention(x, cache=cache)
  
    def testMultiHeadAttention(self):
      attention = transformer.MultiHeadAttention(4, 20)
      queries = tf.random.uniform([4, 5, 10])
      memory = tf.random.uniform([4, 3, 10])
      mask = tf.sequence_mask([1, 3, 2, 2])
      context, _ = attention(queries, memory=memory, mask=mask)
      self.assertListEqual(context.shape.as_list(), [4, 5, 20])
  
    def testMultiHeadAttentionWithCache(self):
      cache = (tf.zeros([4, 4, 0, 5]), tf.zeros([4, 4, 0, 5]))
      attention = transformer.MultiHeadAttention(4, 20)
      memory = tf.random.uniform([4, 3, 10])
      mask = tf.sequence_mask([1, 3, 2, 2])
      x = tf.random.uniform([4, 1, 10])
      y1, cache = attention(x, memory=memory, mask=mask, cache=cache)
      self.assertEqual(cache[0].shape[2], 3)
      self.assertEqual(cache[1].shape[2], 3)
      y2, cache = attention(x, memory=memory, mask=mask, cache=cache)
      self.assertAllEqual(y1, y2)
  
    def testMultiHeadAttentionMask(self):
      attention = transformer.MultiHeadAttention(4, 20, return_attention=True)
      queries = tf.random.uniform([4, 5, 10])
      memory = tf.random.uniform([4, 3, 10])
      mask = tf.sequence_mask([1, 3, 2, 2])
      _, _, attention = attention(queries, memory=memory, mask=mask)
      attention = tf.reshape(attention, [4, -1, 3])
      mask = tf.broadcast_to(tf.expand_dims(mask, 1), attention.shape)
      padding = tf.boolean_mask(attention, tf.logical_not(mask))
      self.assertAllEqual(tf.reduce_sum(padding), 0)
  
  
  if __name__ == "__main__":
    tf.test.main()
  
apis: tensorflow.constant tensorflow.zeros tensorflow.reduce_sum tensorflow.reshape tensorflow.random.uniform tensorflow.test.main tensorflow.expand_dims tensorflow.logical_not numpy.array tensorflow.sequence_mask
api_extract: [(11, 'parameterized.parameterized.expand', 'parameterized.expand', (['[[tf.bool], [tf.float32]]'], {}), False, 'from parameterized import parameterized\n'), (34, 'parameterized.parameterized.expand', 'parameterized.expand', (['[[tf.bool], [tf.float32]]'], {}), False, 'from parameterized import parameterized\n'), (178, 'tensorflow.test.main', 'tf.test.main', ([], {}), True, 'import tensorflow as tf\n'), (69, 'opennmt.layers.transformer.split_heads', 'transformer.split_heads', (['inputs', 'num_heads'], {}), False, 'from opennmt.layers import transformer\n'), (84, 'opennmt.layers.transformer.combine_heads', 'transformer.combine_heads', (['inputs'], {}), False, 'from opennmt.layers import transformer\n'), (98, 'opennmt.layers.transformer.split_heads', 'transformer.split_heads', (['inputs', 'num_heads'], {}), False, 'from opennmt.layers import transformer\n'), (99, 'opennmt.layers.transformer.combine_heads', 'transformer.combine_heads', (['split'], {}), False, 'from opennmt.layers import transformer\n'), (104, 'opennmt.layers.transformer.relative_positions', 'transformer.relative_positions', (['(4)', '(2)'], {}), False, 'from opennmt.layers import transformer\n'), (110, 'opennmt.layers.transformer.FeedForwardNetwork', 'transformer.FeedForwardNetwork', (['(20)', '(10)'], {}), False, 'from opennmt.layers import transformer\n'), (111, 'tensorflow.random.uniform', 'tf.random.uniform', (['[4, 5, 10]'], {}), True, 'import tensorflow as tf\n'), (116, 'opennmt.layers.transformer.MultiHeadAttention', 'transformer.MultiHeadAttention', (['(4)', '(20)'], {}), False, 'from opennmt.layers import transformer\n'), (117, 'tensorflow.random.uniform', 'tf.random.uniform', (['[4, 5, 10]'], {}), True, 'import tensorflow as tf\n'), (118, 'tensorflow.sequence_mask', 'tf.sequence_mask', (['[4, 3, 5, 2]'], {}), True, 'import tensorflow as tf\n'), (124, 'opennmt.layers.transformer.MultiHeadAttention', 'transformer.MultiHeadAttention', (['(4)', '(20)'], {}), False, 'from opennmt.layers import transformer\n'), (125, 'tensorflow.random.uniform', 'tf.random.uniform', (['[4, 1, 10]'], {}), True, 'import tensorflow as tf\n'), (134, 'opennmt.layers.transformer.MultiHeadAttention', 'transformer.MultiHeadAttention', (['(4)', '(20)'], {'maximum_relative_position': '(6)'}), False, 'from opennmt.layers import transformer\n'), (135, 'tensorflow.random.uniform', 'tf.random.uniform', (['[2, 9, 10]'], {}), True, 'import tensorflow as tf\n'), (136, 'tensorflow.sequence_mask', 'tf.sequence_mask', (['[9, 7]'], {}), True, 'import tensorflow as tf\n'), (140, 'opennmt.layers.transformer.MultiHeadAttention', 'transformer.MultiHeadAttention', (['(4)', '(20)'], {'maximum_relative_position': '(6)'}), False, 'from opennmt.layers import transformer\n'), (141, 'tensorflow.random.uniform', 'tf.random.uniform', (['[4, 1, 10]'], {}), True, 'import tensorflow as tf\n'), (146, 'opennmt.layers.transformer.MultiHeadAttention', 'transformer.MultiHeadAttention', (['(4)', '(20)'], {}), False, 'from opennmt.layers import transformer\n'), (147, 'tensorflow.random.uniform', 'tf.random.uniform', (['[4, 5, 10]'], {}), True, 'import tensorflow as tf\n'), (148, 'tensorflow.random.uniform', 'tf.random.uniform', (['[4, 3, 10]'], {}), True, 'import tensorflow as tf\n'), (149, 'tensorflow.sequence_mask', 'tf.sequence_mask', (['[1, 3, 2, 2]'], {}), True, 'import tensorflow as tf\n'), (155, 'opennmt.layers.transformer.MultiHeadAttention', 'transformer.MultiHeadAttention', (['(4)', '(20)'], {}), False, 'from opennmt.layers import transformer\n'), (156, 'tensorflow.random.uniform', 'tf.random.uniform', (['[4, 3, 10]'], {}), True, 'import tensorflow as tf\n'), (157, 'tensorflow.sequence_mask', 'tf.sequence_mask', (['[1, 3, 2, 2]'], {}), True, 'import tensorflow as tf\n'), (158, 'tensorflow.random.uniform', 'tf.random.uniform', (['[4, 1, 10]'], {}), True, 'import tensorflow as tf\n'), (166, 'opennmt.layers.transformer.MultiHeadAttention', 'transformer.MultiHeadAttention', (['(4)', '(20)'], {'return_attention': '(True)'}), False, 'from opennmt.layers import transformer\n'), (167, 'tensorflow.random.uniform', 'tf.random.uniform', (['[4, 5, 10]'], {}), True, 'import tensorflow as tf\n'), (168, 'tensorflow.random.uniform', 'tf.random.uniform', (['[4, 3, 10]'], {}), True, 'import tensorflow as tf\n'), (169, 'tensorflow.sequence_mask', 'tf.sequence_mask', (['[1, 3, 2, 2]'], {}), True, 'import tensorflow as tf\n'), (171, 'tensorflow.reshape', 'tf.reshape', (['attention', '[4, -1, 3]'], {}), True, 'import tensorflow as tf\n'), (28, 'tensorflow.constant', 'tf.constant', (['length'], {}), True, 'import tensorflow as tf\n'), (56, 'tensorflow.constant', 'tf.constant', (['length'], {}), True, 'import tensorflow as tf\n'), (123, 'tensorflow.zeros', 'tf.zeros', (['[4, 4, 0, 5]'], {}), True, 'import tensorflow as tf\n'), (123, 'tensorflow.zeros', 'tf.zeros', (['[4, 4, 0, 5]'], {}), True, 'import tensorflow as tf\n'), (142, 'tensorflow.zeros', 'tf.zeros', (['[4, 4, 0, 5]'], {}), True, 'import tensorflow as tf\n'), (142, 'tensorflow.zeros', 'tf.zeros', (['[4, 4, 0, 5]'], {}), True, 'import tensorflow as tf\n'), (154, 'tensorflow.zeros', 'tf.zeros', (['[4, 4, 0, 5]'], {}), True, 'import tensorflow as tf\n'), (154, 'tensorflow.zeros', 'tf.zeros', (['[4, 4, 0, 5]'], {}), True, 'import tensorflow as tf\n'), (172, 'tensorflow.expand_dims', 'tf.expand_dims', (['mask', '(1)'], {}), True, 'import tensorflow as tf\n'), (173, 'tensorflow.logical_not', 'tf.logical_not', (['mask'], {}), True, 'import tensorflow as tf\n'), (174, 'tensorflow.reduce_sum', 'tf.reduce_sum', (['padding'], {}), True, 'import tensorflow as tf\n'), (14, 'numpy.array', 'np.array', (['[[[1, 0, 0, 0], [1, 1, 0, 0], [1, 1, 0, 0], [1, 1, 0, 0]], [[1, 0, 0, 0], [\n    1, 1, 0, 0], [1, 1, 1, 0], [1, 1, 1, 1]], [[1, 0, 0, 0], [1, 1, 0, 0],\n    [1, 1, 1, 0], [1, 1, 1, 0]]]'], {}), True, 'import numpy as np\n'), (38, 'numpy.array', 'np.array', (['[[[1, 0, 0, 0, 0], [1, 1, 0, 0, 0], [1, 1, 0, 0, 0], [1, 1, 0, 0, 0], [1, 1,\n    0, 0, 0]], [[1, 0, 0, 0, 0], [1, 1, 0, 0, 0], [1, 1, 1, 0, 0], [1, 1, 1,\n    1, 0], [1, 1, 1, 1, 0]], [[1, 0, 0, 0, 0], [1, 1, 0, 0, 0], [1, 1, 1, 0,\n    0], [1, 1, 1, 0, 0], [1, 1, 1, 0, 0]]]'], {}), True, 'import numpy as np\n')]