code: 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,
          is_training=False,
      )
      mesh_impl = mtf.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()
      self.evaluate(init)
      self.evaluate(tf_group)
      actual = self.evaluate(actual_outputs)
  
      self.assertEqual(actual.shape, query.shape)
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  from absl.testing import parameterized
  import mesh_tensorflow
  from mesh_tensorflow import test_utils
  import mock
  import numpy
  import tensorflow.compat.v1
  import tensorflow_probability
  from tensorflow.python.framework import test_util
label_desc: Use image node_id 10 for calling the LayersTest obj's underlying member method code with example usage: obj.testMultiheadAttention(kv_channels, heads) without return types