code: def layout_and_mesh_shape(
      mtf_graph,
      num_machines,
      mtf_outputs=(),
      max_mesh_shape_dimensions=2,
  ):
      """Compute layout rules and mesh shape based on computational graph.
  
      Brute-forces over all possible mesh shapes to find a (layout, mesh_shape)
      pair. Note that the layout optimizer is more efficient when the mesh_shape has
      fewer dimensions, so a smaller max_mesh_shape_dimensions makes this call
      faster.
  
      Args:
        mtf_graph: a mtf.Graph.
        num_machines: integer, a power of two, the number of machines available.
        mtf_outputs: an optional iterable of mtf.Tensor, representing the outputs
            of the computation.
        max_mesh_shape_dimensions: optional integer, the maximum number of
            dimensions to consider in any layout. For example, num_machines=1024 and
            max_mesh_shape_dimensions=2 results in testing the mesh shapes
            "mesh_0:1024", "mesh_0:512;mesh_1:2", "mesh_0:256;mesh_1:4",
            "mesh_0:128;mesh_1:8", "mesh_0:64;mesh_1:16", and "mesh_0:32;mesh_1:32".
            If set to None, there is no maximum.
  
      Returns:
        a (mtf.LayoutRules, mtf.Shape) tuple.
      """
      best_layout_and_mesh_shape = (None, None)
      best_value = None
      for mesh_shape_list in _mesh_shape_iterator(
          num_machines, max_mesh_shape_dimensions
      ):
          mesh_shape = mtf.Shape(
              [
                  mtf.Dimension("mesh_{}".format(i), size)
                  for i, size in enumerate(mesh_shape_list)
              ]
          )
          tf.logging.info(
              "Computing layout for mesh shape: {}".format(mesh_shape)
          )
          estimator = memory_estimator.MemoryEstimator(
              mtf_graph, mesh_shape, mtf_outputs
          )
          optimizer = layout_optimizer.LayoutOptimizer(estimator)
          layout_string = optimizer.solve()
          value = optimizer.evaluate_layout(layout_string)
          if best_value is None or value < best_value:
              best_value = value
              best_layout_and_mesh_shape = (
                  mtf.convert_to_layout_rules(layout_string),
                  mesh_shape,
              )
      return best_layout_and_mesh_shape
  
file_path: mesh/mesh_tensorflow/auto_mtf/api.py
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 tensorflow.compat.v1
label_desc: Use image node_id 2 for calling a global function with example usage: layout_and_mesh_shape(mtf_graph, num_machines, mtf_outputs, max_mesh_shape_dimensions) and returns: best_layout_and_mesh_shape