# 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 six
import tensorflow.compat.v1 as tf
from mesh_tensorflow.auto_mtf import layout_optimizer
from mesh_tensorflow.auto_mtf import memory_estimator

class AssignmentsTest(tf.test.TestCase):

  def testGenerateAssignments(self):
    splittable_dims = {"s1", "s2", "s3"}
    mesh_dims = {"m1": 4, "m2": 8}

    assignments = layout_optimizer._generate_assignments(splittable_dims,
                                                         mesh_dims)
    # Check that some valid assignments of various sizes are included
    self.assertIn({}, assignments)
    self.assertIn({"s3": "m2"}, assignments)
    self.assertIn({"s1": "m2", "s2": "m1"}, assignments)
    # Not allowed to map two splittable dimensions to the same mesh dimension.
    self.assertNotIn({"s1": "m2", "s3": "m2"}, assignments)
    # Check the total number of assignments returned. We are looking for
    # thirteen because one assignment has no entries, six assignments have one
    # entry, and six assignments have two entries.
    self.assertLen(assignments, 13)


class OptimizeLayoutTest(tf.test.TestCase):

  def setUp(self):
    super(OptimizeLayoutTest, self).setUp()
    self.mtf_graph = mtf.Graph()
    self.mesh = mtf.Mesh(self.mtf_graph, "my_mesh")
    self.mesh_shape = mtf.convert_to_shape("m1:4,m2:2")

  def get_layout_optimizer(self):

# 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/layout_optimizer_test.py
import mesh_tensorflow as mtf
import six
import tensorflow.compat.v1 as tf
from mesh_tensorflow.auto_mtf import layout_optimizer
from mesh_tensorflow.auto_mtf import memory_estimator


class AssignmentsTest(tf.test.TestCase):

  def testGenerateAssignments(self):
    splittable_dims = {"s1", "s2", "s3"}
    mesh_dims = {"m1": 4, "m2": 8}

    assignments = layout_optimizer._generate_assignments(splittable_dims,
                                                         mesh_dims)
    # Check that some valid assignments of various sizes are included
    self.assertIn({}, assignments)
    self.assertIn({"s3": "m2"}, assignments)
    self.assertIn({"s1": "m2", "s2": "m1"}, assignments)
    # Not allowed to map two splittable dimensions to the same mesh dimension.
    self.assertNotIn({"s1": "m2", "s3": "m2"}, assignments)
    # Check the total number of assignments returned. We are looking for
    # thirteen because one assignment has no entries, six assignments have one
    # entry, and six assignments have two entries.
    self.assertLen(assignments, 13)


class OptimizeLayoutTest(tf.test.TestCase):

  def setUp(self):
    super(OptimizeLayoutTest, self).setUp()
    self.mtf_graph = mtf.Graph()
    self.mesh = mtf.Mesh(self.mtf_graph, "my_mesh")
    self.mesh_shape = mtf.convert_to_shape("m1:4,m2:2")

  def get_layout_optimizer(self):

    return layout_optimizer.LayoutOptimizer(memory_estimator.MemoryEstimator(