# Path: mesh_tensorflow/auto_mtf/layout_optimizer.py
# class SolverError(Exception):
# class LayoutOptimizer(object):
#   def __init__(self, memory_estimator, scheduler_alg="LIST"):
#   def _preprocess_input(self):
#   def _initialize_variables(self):
#   def _add_constraints(self):
#   def _build_objective_function(self):
#   def _get_memory_contents(self):
#   def solve(self, print_solution=False):
#   def evaluate_layout(self, layout):
# def _global_var_name(splittable_dimension, mesh_dimension):
# def _local_var_name(splittable_dimensions, assignment):
# def _generate_assignments(splittable_dimensions, mesh_dimension_to_size):
# 
# Path: mesh_tensorflow/auto_mtf/memory_estimator.py
# class MemoryEstimator(object):
#   def __init__(self, mtf_graph, mesh_shape, mtf_outputs=()):
#   def get_layout_validator(self):
#   def get_graph_interface(self):
#   def _compute_layout_validator(self):
#   def _compute_graph_interface(self):


import mesh_tensorflow as mtf
import tensorflow.compat.v1 as tf
from mesh_tensorflow.auto_mtf import layout_optimizer
from mesh_tensorflow.auto_mtf import memory_estimator

  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)

# Path: mesh_tensorflow/auto_mtf/layout_optimizer.py
# class SolverError(Exception):
# class LayoutOptimizer(object):
#   def __init__(self, memory_estimator, scheduler_alg="LIST"):
#   def _preprocess_input(self):
#   def _initialize_variables(self):
#   def _add_constraints(self):
#   def _build_objective_function(self):
#   def _get_memory_contents(self):
#   def solve(self, print_solution=False):
#   def evaluate_layout(self, layout):
# def _global_var_name(splittable_dimension, mesh_dimension):
# def _local_var_name(splittable_dimensions, assignment):
# def _generate_assignments(splittable_dimensions, mesh_dimension_to_size):
# 
# Path: mesh_tensorflow/auto_mtf/memory_estimator.py
# class MemoryEstimator(object):
#   def __init__(self, mtf_graph, mesh_shape, mtf_outputs=()):
#   def get_layout_validator(self):
#   def get_graph_interface(self):
#   def _compute_layout_validator(self):
#   def _compute_graph_interface(self):


# Path: mesh_tensorflow/auto_mtf/api.py
import mesh_tensorflow as mtf
import tensorflow.compat.v1 as tf
from mesh_tensorflow.auto_mtf import layout_optimizer
from mesh_tensorflow.auto_mtf import memory_estimator


  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)