code: # coding=utf-8
  # Copyright 2018 The Tensor2Tensor Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  """Tests for Mesh TensorFlow layers."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  from absl.testing import parameterized
  
  from tensor2tensor.layers import common_layers
  from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf
  from tensor2tensor.mesh_tensorflow import mtf_layers
  from tensor2tensor.mesh_tensorflow import placement_mesh_impl
  
  import tensorflow as tf
  
  
  class MtfLayersTest(parameterized.TestCase, tf.test.TestCase):
  
    @parameterized.parameters(
        (4, True),
        (8, False),
    )
    def testDense(self, units, use_bias):
      batch = 2
      channels = 3
      inputs = tf.random_normal([batch, channels])
  
      graph = mtf.Graph()
      mesh = mtf.Mesh(graph, "my_mesh")
      batch_dim = mtf.Dimension("batch", batch)
      channels_dim = mtf.Dimension("channels", channels)
      depth_dim = mtf.Dimension("depth", units)
  
      mtf_inputs = mtf.import_tf_tensor(
          mesh, inputs, shape=mtf.Shape([batch_dim, channels_dim]))
      mtf_outputs = mtf_layers.dense(mtf_inputs,
                                     output_dim=depth_dim,
                                     reduced_dims=[channels_dim],
                                     activation=mtf.relu,
                                     use_bias=use_bias)
      mesh_impl = placement_mesh_impl.PlacementMeshImpl(
          shape=[], layout={}, devices=[""])
      lowering = mtf.Lowering(graph, {mesh: mesh_impl})
      actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
  
      expected_outputs = tf.keras.layers.Dense(units=units,
                                               activation=tf.nn.relu,
                                               use_bias=use_bias)(inputs)
      tf_group = lowering.copy_masters_to_slices()
      init = tf.global_variables_initializer()
      with self.test_session() as sess:
        sess.run(init)
        sess.run(tf_group)
        actual, expected = sess.run([actual_outputs, expected_outputs])
  
      self.assertEqual(actual.shape, expected.shape)
  
    def testLayerNorm(self):
      batch = 2
      channels = 3
      inputs = tf.random_normal([batch, channels])
  
      graph = mtf.Graph()
      mesh = mtf.Mesh(graph, "my_mesh")
      batch_dim = mtf.Dimension("batch", batch)
      channels_dim = mtf.Dimension("channels", channels)
  
      mtf_inputs = mtf.import_tf_tensor(
          mesh, inputs, shape=mtf.Shape([batch_dim, channels_dim]))
      mtf_outputs = mtf_layers.layer_norm(mtf_inputs,
                                          dim=channels_dim)
      mesh_impl = placement_mesh_impl.PlacementMeshImpl(
          shape=[], layout={}, devices=[""])
      lowering = mtf.Lowering(graph, {mesh: mesh_impl})
      actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
  
      expected_outputs = common_layers.layer_norm(inputs)
      tf_group = lowering.copy_masters_to_slices()
      init = tf.global_variables_initializer()
      with self.test_session() as sess:
        sess.run(init)
        sess.run(tf_group)
        actual, expected = sess.run([actual_outputs, expected_outputs])
  
      self.assertEqual(actual.shape, expected.shape)
  
    def testWeightsNonzero(self):
      inputs = tf.constant([[3, 1, 0], [1, 0, 0]])
  
      graph = mtf.Graph()
      mesh = mtf.Mesh(graph, "my_mesh")
      batch_dim = mtf.Dimension("batch", inputs.shape.as_list()[0])
      channels_dim = mtf.Dimension("channels", inputs.shape.as_list()[1])
  
      mtf_inputs = mtf.import_tf_tensor(
          mesh, inputs, shape=mtf.Shape([batch_dim, channels_dim]))
      mtf_outputs = mtf_layers.weights_nonzero(mtf_inputs)
      mesh_impl = placement_mesh_impl.PlacementMeshImpl(
          shape=[], layout={}, devices=[""])
      lowering = mtf.Lowering(graph, {mesh: mesh_impl})
      actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
  
      expected_outputs = common_layers.weights_nonzero(inputs)
      tf_group = lowering.copy_masters_to_slices()
      with self.test_session() as sess:
        sess.run(tf_group)
        actual, expected = sess.run([actual_outputs, expected_outputs])
  
      self.assertAllEqual(actual, expected)
  
    def testDenseReluDense(self):
      batch = 2
      channels = 3
      hidden = 5
      inputs = tf.random_normal([batch, channels])
  
      graph = mtf.Graph()
      mesh = mtf.Mesh(graph, "my_mesh")
      batch_dim = mtf.Dimension("batch", batch)
      channels_dim = mtf.Dimension("channels", channels)
      hidden_dim = mtf.Dimension("hidden", hidden)
  
      mtf_inputs = mtf.import_tf_tensor(
          mesh, inputs, shape=mtf.Shape([batch_dim, channels_dim]))
      mtf_outputs = mtf_layers.dense_relu_dense(mtf_inputs,
                                                hidden_channels=hidden_dim)
      mesh_impl = placement_mesh_impl.PlacementMeshImpl(
          shape=[], layout={}, devices=[""])
      lowering = mtf.Lowering(graph, {mesh: mesh_impl})
      actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
  
      tf_group = lowering.copy_masters_to_slices()
      init = tf.global_variables_initializer()
      with self.test_session() as sess:
        sess.run(init)
        sess.run(tf_group)
        actual = sess.run(actual_outputs)
  
      self.assertEqual(actual.shape, inputs.shape)
  
    @parameterized.parameters(
        (4, 2),
    )
    def testMaskedLocalAttention1D(self, kv_channels, heads):
      batch = 2
      length_q = 16
      length_m = 16
      channels = 3
      query = tf.random_normal([batch, length_q, channels])
      memory = tf.random_normal([batch, length_m, channels])
  
      graph = mtf.Graph()
      mesh = mtf.Mesh(graph, "my_mesh")
      batch_dim = mtf.Dimension("batch", batch)
      length_q_dim = mtf.Dimension("length_q", length_q)
      length_m_dim = mtf.Dimension("length_m", length_m)
      channels_dim = mtf.Dimension("channels", channels)
      kv_channels_dim = mtf.Dimension("kv_channels", kv_channels)
      heads_dim = mtf.Dimension("heads", heads)
  
      mtf_query = mtf.import_tf_tensor(
          mesh, query,
          shape=mtf.Shape([batch_dim, length_q_dim, channels_dim]))
      mtf_memory = mtf.import_tf_tensor(
          mesh, memory,
          shape=mtf.Shape([batch_dim, length_m_dim, channels_dim]))
      mtf_outputs = mtf_layers.masked_local_attention_1d(
          mtf_query,
          mtf_memory,
          kv_channels=kv_channels_dim,
          heads=heads_dim,
          block_length=2)
      mesh_impl = placement_mesh_impl.PlacementMeshImpl(
          shape=[], layout={}, devices=[""])
      lowering = mtf.Lowering(graph, {mesh: mesh_impl})
      actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
  
      tf_group = lowering.copy_masters_to_slices()
      init = tf.global_variables_initializer()
      with self.test_session() as sess:
        sess.run(init)
        sess.run(tf_group)
        actual = sess.run(actual_outputs)
  
      self.assertEqual(actual.shape, (batch, length_q, channels))
  
    @parameterized.parameters(
        (2, 4, 5, 7, 3, 1),
    )
    def testDotProductAttention(
        self, batch, heads, length_q, length_kv, depth_k, depth_v):
      query = tf.random_normal([batch, heads, length_q, depth_k])
      key = tf.random_normal([batch, heads, length_kv, depth_k])
      value = tf.random_normal([batch, heads, length_kv, depth_v])
  
      graph = mtf.Graph()
      mesh = mtf.Mesh(graph, "my_mesh")
      batch_dim = mtf.Dimension("batch", batch)
      heads_dim = mtf.Dimension("heads", heads)
      length_q_dim = mtf.Dimension("length_q", length_q)
      length_kv_dim = mtf.Dimension("length_kv", length_kv)
      depth_k_dim = mtf.Dimension("depth_k", depth_k)
      depth_v_dim = mtf.Dimension("depth_v", depth_v)
  
      mtf_query = mtf.import_tf_tensor(
          mesh, query,
          shape=mtf.Shape(
              [batch_dim, heads_dim, length_q_dim, depth_k_dim]))
      mtf_key = mtf.import_tf_tensor(
          mesh, key,
          shape=mtf.Shape(
              [batch_dim, heads_dim, length_kv_dim, depth_k_dim]))
      mtf_value = mtf.import_tf_tensor(
          mesh, value,
          shape=mtf.Shape(
              [batch_dim, heads_dim, length_kv_dim, depth_v_dim]))
      mtf_outputs = mtf_layers.dot_product_attention(
          mtf_query,
          mtf_key,
          mtf_value,
          mask=None)
      mesh_impl = placement_mesh_impl.PlacementMeshImpl(
          shape=[], layout={}, devices=[""])
      lowering = mtf.Lowering(graph, {mesh: mesh_impl})
      actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
  
      tf_group = lowering.copy_masters_to_slices()
      init = tf.global_variables_initializer()
      with self.test_session() as sess:
        sess.run(init)
        sess.run(tf_group)
        actual = sess.run(actual_outputs)
  
      self.assertEqual(actual.shape, (batch, heads, length_q, depth_v))
  
    @parameterized.parameters(
        (16, 4),
        (32, 8),
    )
    def testMultiheadAttention(self, kv_channels, heads):
      batch = 2
      length = 8
      channels = 3
      query = tf.random_normal([batch, length, channels])
  
      graph = mtf.Graph()
      mesh = mtf.Mesh(graph, "my_mesh")
      batch_dim = mtf.Dimension("batch", batch)
      length_dim = mtf.Dimension("length", length)
      channels_dim = mtf.Dimension("channels", channels)
      kv_channels_dim = mtf.Dimension("kv_channels", kv_channels)
      heads_dim = mtf.Dimension("heads", heads)
  
      mtf_query = mtf.import_tf_tensor(
          mesh, query,
          shape=mtf.Shape([batch_dim, length_dim, channels_dim]))
      mtf_outputs = mtf_layers.multihead_attention(
          mtf_query,
          memory_antecedent=None,
          mask=None,
          kv_channels=kv_channels_dim,
          heads=heads_dim)
      mesh_impl = placement_mesh_impl.PlacementMeshImpl(
          shape=[], layout={}, devices=[""])
      lowering = mtf.Lowering(graph, {mesh: mesh_impl})
      actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
  
      tf_group = lowering.copy_masters_to_slices()
      init = tf.global_variables_initializer()
      with self.test_session() as sess:
        sess.run(init)
        sess.run(tf_group)
        actual = sess.run(actual_outputs)
  
      self.assertEqual(actual.shape, query.shape)
  
  if __name__ == "__main__":
    tf.test.main()
  
apis: tensorflow.constant tensorflow.keras.layers.Dense tensorflow.test.main tensorflow.global_variables_initializer tensorflow.random_normal
api_extract: [(33, 'absl.testing.parameterized.parameters', 'parameterized.parameters', (['(4, True)', '(8, False)'], {}), False, 'from absl.testing import parameterized\n'), (155, 'absl.testing.parameterized.parameters', 'parameterized.parameters', (['(4, 2)'], {}), False, 'from absl.testing import parameterized\n'), (201, 'absl.testing.parameterized.parameters', 'parameterized.parameters', (['(2, 4, 5, 7, 3, 1)'], {}), False, 'from absl.testing import parameterized\n'), (250, 'absl.testing.parameterized.parameters', 'parameterized.parameters', (['(16, 4)', '(32, 8)'], {}), False, 'from absl.testing import parameterized\n'), (292, 'tensorflow.test.main', 'tf.test.main', ([], {}), True, 'import tensorflow as tf\n'), (40, 'tensorflow.random_normal', 'tf.random_normal', (['[batch, channels]'], {}), True, 'import tensorflow as tf\n'), (42, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Graph', 'mtf.Graph', ([], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (43, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Mesh', 'mtf.Mesh', (['graph', '"""my_mesh"""'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (44, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""batch"""', 'batch'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (45, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""channels"""', 'channels'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (46, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""depth"""', 'units'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (50, 'tensor2tensor.mesh_tensorflow.mtf_layers.dense', 'mtf_layers.dense', (['mtf_inputs'], {'output_dim': 'depth_dim', 'reduced_dims': '[channels_dim]', 'activation': 'mtf.relu', 'use_bias': 'use_bias'}), False, 'from tensor2tensor.mesh_tensorflow import mtf_layers\n'), (55, 'tensor2tensor.mesh_tensorflow.placement_mesh_impl.PlacementMeshImpl', 'placement_mesh_impl.PlacementMeshImpl', ([], {'shape': '[]', 'layout': '{}', 'devices': "['']"}), False, 'from tensor2tensor.mesh_tensorflow import placement_mesh_impl\n'), (57, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Lowering', 'mtf.Lowering', (['graph', '{mesh: mesh_impl}'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (64, 'tensorflow.global_variables_initializer', 'tf.global_variables_initializer', ([], {}), True, 'import tensorflow as tf\n'), (75, 'tensorflow.random_normal', 'tf.random_normal', (['[batch, channels]'], {}), True, 'import tensorflow as tf\n'), (77, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Graph', 'mtf.Graph', ([], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (78, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Mesh', 'mtf.Mesh', (['graph', '"""my_mesh"""'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (79, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""batch"""', 'batch'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (80, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""channels"""', 'channels'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (84, 'tensor2tensor.mesh_tensorflow.mtf_layers.layer_norm', 'mtf_layers.layer_norm', (['mtf_inputs'], {'dim': 'channels_dim'}), False, 'from tensor2tensor.mesh_tensorflow import mtf_layers\n'), (86, 'tensor2tensor.mesh_tensorflow.placement_mesh_impl.PlacementMeshImpl', 'placement_mesh_impl.PlacementMeshImpl', ([], {'shape': '[]', 'layout': '{}', 'devices': "['']"}), False, 'from tensor2tensor.mesh_tensorflow import placement_mesh_impl\n'), (88, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Lowering', 'mtf.Lowering', (['graph', '{mesh: mesh_impl}'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (91, 'tensor2tensor.layers.common_layers.layer_norm', 'common_layers.layer_norm', (['inputs'], {}), False, 'from tensor2tensor.layers import common_layers\n'), (93, 'tensorflow.global_variables_initializer', 'tf.global_variables_initializer', ([], {}), True, 'import tensorflow as tf\n'), (102, 'tensorflow.constant', 'tf.constant', (['[[3, 1, 0], [1, 0, 0]]'], {}), True, 'import tensorflow as tf\n'), (104, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Graph', 'mtf.Graph', ([], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (105, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Mesh', 'mtf.Mesh', (['graph', '"""my_mesh"""'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (111, 'tensor2tensor.mesh_tensorflow.mtf_layers.weights_nonzero', 'mtf_layers.weights_nonzero', (['mtf_inputs'], {}), False, 'from tensor2tensor.mesh_tensorflow import mtf_layers\n'), (112, 'tensor2tensor.mesh_tensorflow.placement_mesh_impl.PlacementMeshImpl', 'placement_mesh_impl.PlacementMeshImpl', ([], {'shape': '[]', 'layout': '{}', 'devices': "['']"}), False, 'from tensor2tensor.mesh_tensorflow import placement_mesh_impl\n'), (114, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Lowering', 'mtf.Lowering', (['graph', '{mesh: mesh_impl}'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (117, 'tensor2tensor.layers.common_layers.weights_nonzero', 'common_layers.weights_nonzero', (['inputs'], {}), False, 'from tensor2tensor.layers import common_layers\n'), (129, 'tensorflow.random_normal', 'tf.random_normal', (['[batch, channels]'], {}), True, 'import tensorflow as tf\n'), (131, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Graph', 'mtf.Graph', ([], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (132, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Mesh', 'mtf.Mesh', (['graph', '"""my_mesh"""'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (133, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""batch"""', 'batch'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (134, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""channels"""', 'channels'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (135, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""hidden"""', 'hidden'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (139, 'tensor2tensor.mesh_tensorflow.mtf_layers.dense_relu_dense', 'mtf_layers.dense_relu_dense', (['mtf_inputs'], {'hidden_channels': 'hidden_dim'}), False, 'from tensor2tensor.mesh_tensorflow import mtf_layers\n'), (141, 'tensor2tensor.mesh_tensorflow.placement_mesh_impl.PlacementMeshImpl', 'placement_mesh_impl.PlacementMeshImpl', ([], {'shape': '[]', 'layout': '{}', 'devices': "['']"}), False, 'from tensor2tensor.mesh_tensorflow import placement_mesh_impl\n'), (143, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Lowering', 'mtf.Lowering', (['graph', '{mesh: mesh_impl}'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (147, 'tensorflow.global_variables_initializer', 'tf.global_variables_initializer', ([], {}), True, 'import tensorflow as tf\n'), (163, 'tensorflow.random_normal', 'tf.random_normal', (['[batch, length_q, channels]'], {}), True, 'import tensorflow as tf\n'), (164, 'tensorflow.random_normal', 'tf.random_normal', (['[batch, length_m, channels]'], {}), True, 'import tensorflow as tf\n'), (166, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Graph', 'mtf.Graph', ([], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (167, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Mesh', 'mtf.Mesh', (['graph', '"""my_mesh"""'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (168, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""batch"""', 'batch'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (169, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""length_q"""', 'length_q'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (170, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""length_m"""', 'length_m'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (171, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""channels"""', 'channels'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (172, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""kv_channels"""', 'kv_channels'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (173, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""heads"""', 'heads'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (181, 'tensor2tensor.mesh_tensorflow.mtf_layers.masked_local_attention_1d', 'mtf_layers.masked_local_attention_1d', (['mtf_query', 'mtf_memory'], {'kv_channels': 'kv_channels_dim', 'heads': 'heads_dim', 'block_length': '(2)'}), False, 'from tensor2tensor.mesh_tensorflow import mtf_layers\n'), (187, 'tensor2tensor.mesh_tensorflow.placement_mesh_impl.PlacementMeshImpl', 'placement_mesh_impl.PlacementMeshImpl', ([], {'shape': '[]', 'layout': '{}', 'devices': "['']"}), False, 'from tensor2tensor.mesh_tensorflow import placement_mesh_impl\n'), (189, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Lowering', 'mtf.Lowering', (['graph', '{mesh: mesh_impl}'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (193, 'tensorflow.global_variables_initializer', 'tf.global_variables_initializer', ([], {}), True, 'import tensorflow as tf\n'), (206, 'tensorflow.random_normal', 'tf.random_normal', (['[batch, heads, length_q, depth_k]'], {}), True, 'import tensorflow as tf\n'), (207, 'tensorflow.random_normal', 'tf.random_normal', (['[batch, heads, length_kv, depth_k]'], {}), True, 'import tensorflow as tf\n'), (208, 'tensorflow.random_normal', 'tf.random_normal', (['[batch, heads, length_kv, depth_v]'], {}), True, 'import tensorflow as tf\n'), (210, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Graph', 'mtf.Graph', ([], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (211, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Mesh', 'mtf.Mesh', (['graph', '"""my_mesh"""'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (212, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""batch"""', 'batch'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (213, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""heads"""', 'heads'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (214, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""length_q"""', 'length_q'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (215, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""length_kv"""', 'length_kv'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (216, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""depth_k"""', 'depth_k'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (217, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""depth_v"""', 'depth_v'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (231, 'tensor2tensor.mesh_tensorflow.mtf_layers.dot_product_attention', 'mtf_layers.dot_product_attention', (['mtf_query', 'mtf_key', 'mtf_value'], {'mask': 'None'}), False, 'from tensor2tensor.mesh_tensorflow import mtf_layers\n'), (236, 'tensor2tensor.mesh_tensorflow.placement_mesh_impl.PlacementMeshImpl', 'placement_mesh_impl.PlacementMeshImpl', ([], {'shape': '[]', 'layout': '{}', 'devices': "['']"}), False, 'from tensor2tensor.mesh_tensorflow import placement_mesh_impl\n'), (238, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Lowering', 'mtf.Lowering', (['graph', '{mesh: mesh_impl}'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (242, 'tensorflow.global_variables_initializer', 'tf.global_variables_initializer', ([], {}), True, 'import tensorflow as tf\n'), (258, 'tensorflow.random_normal', 'tf.random_normal', (['[batch, length, channels]'], {}), True, 'import tensorflow as tf\n'), (260, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Graph', 'mtf.Graph', ([], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (261, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Mesh', 'mtf.Mesh', (['graph', '"""my_mesh"""'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (262, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""batch"""', 'batch'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (263, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""length"""', 'length'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (264, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""channels"""', 'channels'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (265, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""kv_channels"""', 'kv_channels'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (266, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Dimension', 'mtf.Dimension', (['"""heads"""', 'heads'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (271, 'tensor2tensor.mesh_tensorflow.mtf_layers.multihead_attention', 'mtf_layers.multihead_attention', (['mtf_query'], {'memory_antecedent': 'None', 'mask': 'None', 'kv_channels': 'kv_channels_dim', 'heads': 'heads_dim'}), False, 'from tensor2tensor.mesh_tensorflow import mtf_layers\n'), (277, 'tensor2tensor.mesh_tensorflow.placement_mesh_impl.PlacementMeshImpl', 'placement_mesh_impl.PlacementMeshImpl', ([], {'shape': '[]', 'layout': '{}', 'devices': "['']"}), False, 'from tensor2tensor.mesh_tensorflow import placement_mesh_impl\n'), (279, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Lowering', 'mtf.Lowering', (['graph', '{mesh: mesh_impl}'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (283, 'tensorflow.global_variables_initializer', 'tf.global_variables_initializer', ([], {}), True, 'import tensorflow as tf\n'), (60, 'tensorflow.keras.layers.Dense', 'tf.keras.layers.Dense', ([], {'units': 'units', 'activation': 'tf.nn.relu', 'use_bias': 'use_bias'}), True, 'import tensorflow as tf\n'), (49, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Shape', 'mtf.Shape', (['[batch_dim, channels_dim]'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (83, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Shape', 'mtf.Shape', (['[batch_dim, channels_dim]'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (110, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Shape', 'mtf.Shape', (['[batch_dim, channels_dim]'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (138, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Shape', 'mtf.Shape', (['[batch_dim, channels_dim]'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (177, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Shape', 'mtf.Shape', (['[batch_dim, length_q_dim, channels_dim]'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (180, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Shape', 'mtf.Shape', (['[batch_dim, length_m_dim, channels_dim]'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (221, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Shape', 'mtf.Shape', (['[batch_dim, heads_dim, length_q_dim, depth_k_dim]'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (225, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Shape', 'mtf.Shape', (['[batch_dim, heads_dim, length_kv_dim, depth_k_dim]'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (229, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Shape', 'mtf.Shape', (['[batch_dim, heads_dim, length_kv_dim, depth_v_dim]'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n'), (270, 'tensor2tensor.mesh_tensorflow.mesh_tensorflow.Shape', 'mtf.Shape', (['[batch_dim, length_dim, channels_dim]'], {}), True, 'from tensor2tensor.mesh_tensorflow import mesh_tensorflow as mtf\n')]