# Path: mesh_tensorflow/auto_mtf/valid_layouts.py
# class LayoutValidator(object):
#   def __init__(self, mtf_graph, mesh_shape):
#   def splittable_mtf_dimension_names(self):
#   def mesh_dimension_name_to_size(self):
#   def is_valid_assignment(self, mtf_dimension_name, mesh_dimension_name):
#   def _initialize_splittable_dimensions(self, mtf_graph):
#   def _initialize_mtf_dimension_name_to_size_gcd(self, mtf_graph):
#   def _initialize_mesh_dimension_name_to_size(self, mesh_shape):


import mesh_tensorflow as mtf
import tensorflow.compat.v1 as tf
from mesh_tensorflow.auto_mtf import valid_layouts
    mesh = mtf.Mesh(graph, "my_mesh")

    a_dim = mtf.Dimension("a", 5)
    b_dim = mtf.Dimension("b", 10)
    concat_dim1 = mtf.Dimension("concat", 15)
    concat_dim2 = mtf.Dimension("concat", 20)

    x1 = mtf.zeros(mesh, mtf.Shape([a_dim, b_dim, concat_dim1]))
    x2 = mtf.zeros(mesh, mtf.Shape([a_dim, b_dim, concat_dim2]))
    mtf.ConcatOperation([x1, x2], "concat")

    # We add a tensor with anonymous shape, which is supposed to be
    # unsplittable (i.e. none of its dimensions show up during
    # test_SplittableMtfDimensionNames).
    _ = mtf.zeros(mesh, mtf.anonymous_shape(mtf.Shape([a_dim, b_dim])))

    mesh_shape = mtf.Shape([("m1", 4), ("m2", 2)])
    self.valid_layouts = valid_layouts.LayoutValidator(graph, mesh_shape)

  def test_SplittableMtfDimensionNames(self):
    self.assertEqual(self.valid_layouts.splittable_mtf_dimension_names,
                     set(["a", "b"]))

  def test_MeshDimensionNameToSize(self):
    self.assertEqual(self.valid_layouts.mesh_dimension_name_to_size,
                     {"m1": 4, "m2": 2})

  def test_is_valid_assignment(self):
    # Due to divisibility, the a dimension cannot be assigned to m1 or m2.
    self.assertFalse(self.valid_layouts.is_valid_assignment("a", "m1"))

# Path: mesh_tensorflow/auto_mtf/valid_layouts.py
# class LayoutValidator(object):
#   def __init__(self, mtf_graph, mesh_shape):
#   def splittable_mtf_dimension_names(self):
#   def mesh_dimension_name_to_size(self):
#   def is_valid_assignment(self, mtf_dimension_name, mesh_dimension_name):
#   def _initialize_splittable_dimensions(self, mtf_graph):
#   def _initialize_mtf_dimension_name_to_size_gcd(self, mtf_graph):
#   def _initialize_mesh_dimension_name_to_size(self, mesh_shape):


# Path: mesh_tensorflow/auto_mtf/valid_layouts_test.py
import mesh_tensorflow as mtf
import tensorflow.compat.v1 as tf
from mesh_tensorflow.auto_mtf import valid_layouts

    mesh = mtf.Mesh(graph, "my_mesh")

    a_dim = mtf.Dimension("a", 5)
    b_dim = mtf.Dimension("b", 10)
    concat_dim1 = mtf.Dimension("concat", 15)
    concat_dim2 = mtf.Dimension("concat", 20)

    x1 = mtf.zeros(mesh, mtf.Shape([a_dim, b_dim, concat_dim1]))
    x2 = mtf.zeros(mesh, mtf.Shape([a_dim, b_dim, concat_dim2]))
    mtf.ConcatOperation([x1, x2], "concat")

    # We add a tensor with anonymous shape, which is supposed to be
    # unsplittable (i.e. none of its dimensions show up during
    # test_SplittableMtfDimensionNames).
    _ = mtf.zeros(mesh, mtf.anonymous_shape(mtf.Shape([a_dim, b_dim])))

    mesh_shape = mtf.Shape([("m1", 4), ("m2", 2)])
    self.valid_layouts = valid_layouts.LayoutValidator(graph, mesh_shape)

  def test_SplittableMtfDimensionNames(self):
    self.assertEqual(self.valid_layouts.splittable_mtf_dimension_names,
                     set(["a", "b"]))

  def test_MeshDimensionNameToSize(self):
    self.assertEqual(self.valid_layouts.mesh_dimension_name_to_size,
                     {"m1": 4, "m2": 2})

  def test_is_valid_assignment(self):
    # Due to divisibility, the a dimension cannot be assigned to m1 or m2.
    self.assertFalse(self.valid_layouts.is_valid_assignment("a", "m1"))

    self.assertFalse(self.valid_layouts.is_valid_assignment("a", "m2"))