code: def testLayout(self):
      # Construct a Mesh TensorFlow graph and mesh.
      mtf_graph = mtf.Graph()
      mesh = mtf.Mesh(mtf_graph, "my_mesh")
      x = mtf.zeros(mesh, "a:10,b:5")
      y = mtf.zeros(mesh, "b:5,c:20")
      z = mtf.einsum([x, y], "a:10,c:20")
  
      # Decide on a mesh shape.
      mesh_shape = mtf.convert_to_shape("m1:4,m2:2")
  
      # Compute a layout based on the graph and mesh.
      # Note that knowing the identity of the outputs is important to the
      # optimization since they cannot be freed.
      layout = mtf.auto_mtf.layout(mtf_graph, mesh_shape, [z])
  
      a_dim = mtf.convert_to_dimension(("a", 10))
      b_dim = mtf.convert_to_dimension(("b", 5))
      c_dim = mtf.convert_to_dimension(("c", 20))
  
      self.assertEqual(
          layout.tensor_dimension_to_mesh_axis(a_dim, mesh_shape), 1
      )
      self.assertIsNone(
          layout.tensor_dimension_to_mesh_axis(b_dim, mesh_shape)
      )
      self.assertEqual(
          layout.tensor_dimension_to_mesh_axis(c_dim, mesh_shape), 0
      )
  
file_path: mesh/mesh_tensorflow/auto_mtf/api_test.py
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import mesh_tensorflow
  import mesh_tensorflow.auto_mtf
  import mesh_tensorflow.auto_mtf.api
  import tensorflow.compat.v1
label_desc: Use image node_id 1 for calling the LayoutTest obj's underlying member method code with example usage: obj.testLayout() without return types