code: def testOptimizeLayout(self):
      x1 = mtf.zeros(self.mesh, "a:10,b:5")
      x2 = mtf.zeros(self.mesh, "b:5,c:20")
      mtf.einsum([x1, x2], "a:10,c:20")
      optimizer = self.get_layout_optimizer()
  
      # Cut dimensions to make them equally sized.
      layout = optimizer.solve()
      self.assertEqual(layout, "a:m2;c:m1")
  
      # This optimal layout should have the lowest value.
      layout_value = optimizer.evaluate_layout(layout)
      self.assertLessEqual(
          layout_value, optimizer.evaluate_layout("a:m1;b:m2")
      )
      self.assertLessEqual(
          layout_value, optimizer.evaluate_layout("a:m1;c:m2")
      )
      self.assertLessEqual(
          layout_value, optimizer.evaluate_layout("b:m1;a:m2")
      )
      self.assertLessEqual(
          layout_value, optimizer.evaluate_layout("b:m1;c:m2")
      )
      self.assertLessEqual(
          layout_value, optimizer.evaluate_layout("c:m1;b:m2")
      )
      self.assertEqual(
          layout_value, optimizer.evaluate_layout("c:m1;a:m2")
      )
  
imports: from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  import mesh_tensorflow
  from mesh_tensorflow.auto_mtf import layout_optimizer
  from mesh_tensorflow.auto_mtf import memory_estimator
  import six
  import tensorflow.compat.v1
label_desc: Use image node_id 3 for calling the OptimizeLayoutTest obj's underlying member method code with example usage: obj.testOptimizeLayout() without return types