text: # coding=utf-8
  # Copyright 2019 The Mesh TensorFlow Authors.
  #
  # Licensed under the Apache License, Version 2.0 (the "License");
  # you may not use this file except in compliance with the License.
  # You may obtain a copy of the License at
  #
  #     http://www.apache.org/licenses/LICENSE-2.0
  #
  # Unless required by applicable law or agreed to in writing, software
  # distributed under the License is distributed on an "AS IS" BASIS,
  # WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
  # See the License for the specific language governing permissions and
  # limitations under the License.
  
  """Tests for Mesh TensorFlow."""
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  from absl.testing import parameterized
  
  import mesh_tensorflow as mtf
  import tensorflow.compat.v1 as tf
  from tensorflow.python.framework import test_util  # pylint:disable=g-direct-tensorflow-import
  
  
  class LaidOutTensor(object):
    """LaidOutTensor (see placement_mesh_impl.py, simd_mesh_impl.py) for tests."""
  
    def __init__(self, tensor_list):
      self.tensor_list = tensor_list
  
  
  class MeshTensorFlowTest(parameterized.TestCase, tf.test.TestCase):
  
    @parameterized.parameters(
        (mtf.Dimension("x", 5),),
        (("x", 5),),
    )
    def testConvertToDimension(self, inputs):
      dimension = mtf.convert_to_dimension(inputs)
      self.assertEqual(dimension.name, "x")
      self.assertEqual(dimension.size, 5)
  
    def testConvertToDimensionGenericInputs(self):
      dimension = mtf.convert_to_dimension(None)
      self.assertEqual(dimension, None)
      with self.assertRaises(TypeError):
        mtf.convert_to_dimension(5)
  
    @parameterized.parameters(
        (mtf.Shape([mtf.Dimension("x", 4),
                    mtf.Dimension("y", 8)]),),
        ("x:4;y:8",),
        ("x:4.y:8",),
        ("x:4 y:8",),
        ("x:4,y:8",),
    )
    def testConvertToShape(self, inputs):
      shape = mtf.convert_to_shape(inputs)
      self.assertEqual(shape, mtf.Shape([mtf.Dimension("x", 4),
                                         mtf.Dimension("y", 8)]))
  
    def testConvertToShapeGenericInputs(self):
      shape = mtf.convert_to_shape([])
      self.assertEqual(shape.dims, [])
      shape = mtf.convert_to_shape(None)
      self.assertEqual(shape, None)
      with self.assertRaises(ValueError):
        mtf.convert_to_shape("x;4")
  
    @parameterized.parameters(
        (mtf.LayoutRules([("d_ff", "model"), ("heads", "model")]),),
        ("d_ff:model;heads:model",),
        ("d_ff:model.heads:model",),
        ("d_ff:model heads:model",),
        ("d_ff:model,heads:model",),
        ([("d_ff", "model"), ("heads", "model")],),
    )
    def testConvertToLayoutRules(self, inputs):
      layout_rules = mtf.convert_to_layout_rules(inputs)
      self.assertEqual(
          layout_rules._pairs,
          mtf.LayoutRules([("d_ff", "model"), ("heads", "model")])._pairs)
  
    def testConvertToLayoutRulesGenericInputs(self):
      with self.assertRaises(ValueError):
        mtf.convert_to_layout_rules("d_ff;heads")
  
    def testTensorLayout(self):
      tensor_layout = mtf.TensorLayout([0, 2, 1])
      self.assertEqual(tensor_layout.mesh_axis_to_tensor_axis(0), ())
      self.assertEqual(tensor_layout.mesh_axis_to_tensor_axis(1), (0,))
      self.assertEqual(tensor_layout.mesh_axis_to_tensor_axis(2), (0, 2))
      tensor_layout = mtf.TensorLayout([None, 0])
      self.assertFalse(tensor_layout.is_fully_replicated)
      tensor_layout = mtf.TensorLayout([None, None, None])
      self.assertTrue(tensor_layout.is_fully_replicated)
  
    def testGraph(self):
      graph = mtf.Graph()
      self.assertEmpty(graph.operations)
      self.assertEmpty(graph.trainable_variables)
      self.assertEmpty(graph.all_variables)
      mesh = mtf.Mesh(graph, "mesh_test")
      _ = mtf.import_tf_tensor(mesh,
                               tf_tensor=tf.constant(0.),
                               shape=mtf.Shape([]))
      self.assertLen(graph.operations, 1)
      self.assertEmpty(graph.trainable_variables)
      self.assertEmpty(graph.all_variables)
      _ = mtf.get_variable(mesh, "variable_0", mtf.Shape([]), trainable=True)
      self.assertLen(graph.operations, 2)
      self.assertLen(graph.trainable_variables, 1)
      self.assertLen(graph.all_variables, 1)
      _ = mtf.get_variable(mesh, "variable_1", mtf.Shape([]), trainable=False)
      self.assertLen(graph.operations, 3)
      self.assertLen(graph.trainable_variables, 1)
      self.assertLen(graph.all_variables, 2)
  
    def testGraphNames(self):
      # Standard Usage.
      graph = mtf.Graph()
      self.assertEqual(graph.unique_name("a"), "a")
      self.assertEqual(graph.unique_name("a"), "a_1")
      self.assertEqual(graph.unique_name("a"), "a_2")
  
      # Edge cases, the user may choose the name "a_1".
      graph = mtf.Graph()
      self.assertEqual(graph.unique_name("a"), "a")
      self.assertEqual(graph.unique_name("a"), "a_1")
      self.assertEqual(graph.unique_name("a_1"), "a_1_1")
  
      graph = mtf.Graph()
      self.assertEqual(graph.unique_name("a"), "a")
      self.assertEqual(graph.unique_name("a_1"), "a_1")
      self.assertEqual(graph.unique_name("a"), "a_2")
  
      # Case insensitive.
      graph = mtf.Graph()
      self.assertEqual(graph.unique_name("a"), "a")
      self.assertEqual(graph.unique_name("A"), "A_1")
  
    @test_util.run_in_graph_and_eager_modes()
    def testLowering(self):
      graph = mtf.Graph()
      mesh = mtf.Mesh(graph, "my_mesh")
      inputs = tf.constant(0.)
      mtf_inputs = mtf.import_tf_tensor(mesh,
                                        tf_tensor=inputs,
                                        shape=mtf.Shape([]))
      mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl(
          shape=[], layout={}, devices=[""])
      lowering = mtf.Lowering(graph, {mesh: mesh_impl})
  
      outputs = lowering.export_to_tf_tensor(mtf_inputs)
      inputs_value, outputs_value = self.evaluate([inputs, outputs])
      self.assertEqual(inputs_value, outputs_value)
  
      # Check that methods run without error.
      _ = lowering.copy_masters_to_slices()
      _ = lowering.copy_slices_to_masters()
  
    def testMesh(self):
      graph = mtf.Graph()
      mesh = mtf.Mesh(graph, "my_mesh")
      self.assertEqual(mesh.graph, graph)
  
    def testMeshImpl(self):
      shape = mtf.Shape([mtf.Dimension("batch", 4),
                         mtf.Dimension("model", 8)])
      layout_rules = mtf.LayoutRules([("batch", "batch"),
                                      ("d_ff", "model"),
                                      ("heads", "model")])
      mesh_impl = mtf.MeshImpl(shape=shape, layout_rules=layout_rules)
      self.assertEqual(mesh_impl.shape, shape)
      self.assertLen(shape, mesh_impl.ndims)
      self.assertEqual(mesh_impl.layout_rules, layout_rules)
      self.assertEqual(mesh_impl.size, shape.size)
      self.assertTrue(mesh_impl.supports_control_dependencies)
  
      batch = mtf.Dimension("batch", 128)
      length = mtf.Dimension("length", 500)
      d_ff = mtf.Dimension("d_ff", 2048)
      heads = mtf.Dimension("heads", 8)
      self.assertEqual(mesh_impl.tensor_dimension_to_mesh_axis(batch), 0)
      self.assertEqual(mesh_impl.tensor_dimension_to_mesh_axis(d_ff), 1)
      self.assertEqual(mesh_impl.tensor_dimension_to_mesh_axis(heads), 1)
      self.assertEqual(mesh_impl.tensor_layout(mtf.Shape([batch, length, d_ff])),
                       mtf.TensorLayout([0, None, 1]))
  
  
  class OperationSplittabilityTest(tf.test.TestCase):
  
    def setUp(self):
      super(OperationSplittabilityTest, self).setUp()
      self.graph = mtf.Graph()
      self.mesh = mtf.Mesh(self.graph, "my_mesh")
  
      self.a_dim = mtf.Dimension("a", 5)
      self.b_dim = mtf.Dimension("b", 10)
      self.c_dim = mtf.Dimension("c", 15)
  
      self.ab_shape = mtf.Shape([self.a_dim, self.b_dim])
      self.x = mtf.zeros(self.mesh, self.ab_shape)
  
      self.batch_dim = mtf.Dimension("batch", 100)
      self.grid_h_dim = mtf.Dimension("grid_h", 10)
      self.grid_w_dim = mtf.Dimension("grid_w", 10)
      self.filter_h_dim = mtf.Dimension("filter_h", 5)
      self.filter_w_dim = mtf.Dimension("filter_w", 5)
      self.in_dim = mtf.Dimension("in", 10)
      self.out_dim = mtf.Dimension("out", 10)
      self.image = mtf.zeros(self.mesh, [self.batch_dim, self.grid_h_dim,
                                         self.grid_w_dim, self.in_dim])
  
    def testOperation(self):
      operation = mtf.Operation([self.x], name="operation")
  
      # Everything is splittable.
      self.assertEqual(
          operation._initialize_all_dimensions_as_splittable(),
          (frozenset(["a", "b"]), frozenset()))
  
      # Everything is unsplittable.
      self.assertEqual(
          operation._initialize_splittable_and_unsplittable_dims("unsplittable"),
          (frozenset(), frozenset(["a", "b"])))
  
      # Everything is unsplittable except dimension "b".
      self.assertEqual(
          operation._initialize_splittable_and_unsplittable_dims(
              "unsplittable", ["b"]),
          (frozenset(["b"]), frozenset(["a"])))
  
      self.assertRaises(
          ValueError,
          operation._initialize_splittable_and_unsplittable_dims,
          "invalid")
  
    def testSlicewiseOperationAndGenericGradOperation(self):
      slicewise_operation = mtf.SlicewiseOperation(
          tf.exp,
          [self.x],
          [self.x.shape],
          [self.x.dtype],
          splittable_dims=[self.a_dim],  # pretend only dim "a" can be split.
          grad_function=lambda op, dy: [dy * op.outputs[0]],
          name="component-wise exp")
  
      self.assertEqual(slicewise_operation.splittable_dims, frozenset(["a"]))
      self.assertEqual(slicewise_operation.unsplittable_dims, frozenset(["b"]))
  
      generic_grad_operation = mtf.GenericGradOperation(slicewise_operation,
                                                        [self.x])
  
      self.assertEqual(generic_grad_operation.splittable_dims,
                       frozenset(["a", "b"]))
      self.assertEqual(generic_grad_operation.unsplittable_dims,
                       frozenset())
  
    def testScalarMultiplyOperationandScalarAddOperation(self):
      scalar = 2.0
      scalar_multiply_operation = mtf.ScalarMultiplyOperation(self.x, scalar)
      self.assertEqual(scalar_multiply_operation.splittable_dims,
                       frozenset(["a", "b"]))
      self.assertEqual(scalar_multiply_operation.unsplittable_dims, frozenset())
  
      scalar_add_operation = mtf.ScalarAddOperation(self.x, scalar)
      self.assertEqual(scalar_add_operation.splittable_dims,
                       frozenset(["a", "b"]))
      self.assertEqual(scalar_add_operation.unsplittable_dims, frozenset())
  
    def testBinaryOpWithBroadcasting(self):
      x2 = mtf.zeros(self.mesh, mtf.Shape([self.a_dim, self.c_dim]))
      binary_op_with_broadcasting = mtf.BinaryOpWithBroadcasting(
          tf.less,
          self.x,
          x2,
          mtf.Shape([self.a_dim, self.b_dim, self.c_dim]),
          tf.bool,
          name="less with broadcasting")
  
      self.assertEqual(binary_op_with_broadcasting.splittable_dims,
                       frozenset(["a", "b", "c"]))
      self.assertEqual(binary_op_with_broadcasting.unsplittable_dims, frozenset())
  
    def testBroadcastOperation(self):
      broadcast_operation = mtf.BroadcastOperation(
          self.x, mtf.Shape([self.b_dim, self.c_dim, self.a_dim]))
      self.assertEqual(broadcast_operation.splittable_dims,
                       frozenset(["a", "b", "c"]))
      self.assertEqual(broadcast_operation.unsplittable_dims, frozenset())
  
    def testReduceOperation(self):
      reduce_operation = mtf.ReduceOperation(self.x, mtf.Shape([self.b_dim]),
                                             "sum")
      self.assertEqual(reduce_operation.splittable_dims, frozenset(["a", "b"]))
      self.assertEqual(reduce_operation.unsplittable_dims, frozenset())
  
    def testPoolOperation(self):
      reduce_operation = mtf.PoolOperation(self.image, [2, 2], [2, 2], "AVG_2D")
      self.assertEqual(reduce_operation.splittable_dims,
                       frozenset(["batch", "in"]))
      self.assertEqual(reduce_operation.unsplittable_dims,
                       frozenset(["grid_h", "grid_w"]))
  
    def testConcatOperation(self):
      concat_dim1 = mtf.Dimension("concat", 5)
      concat_dim2 = mtf.Dimension("concat", 7)
  
      x1 = mtf.zeros(self.mesh, mtf.Shape([self.a_dim, self.b_dim, concat_dim1]))
      x2 = mtf.zeros(self.mesh, mtf.Shape([self.a_dim, self.b_dim, concat_dim2]))
  
      concat_operation = mtf.ConcatOperation([x1, x2], "concat")
      self.assertEqual(concat_operation.splittable_dims, frozenset(["a", "b"]))
      self.assertEqual(concat_operation.unsplittable_dims, frozenset(["concat"]))
  
    def testSplitOperation(self):
      split_operation = mtf.SplitOperation(self.x, self.b_dim, [3, 7])
      self.assertEqual(split_operation.splittable_dims, frozenset(["a"]))
      self.assertEqual(split_operation.unsplittable_dims, frozenset(["b"]))
  
    def testStackOperation(self):
      stack_operation = mtf.StackOperation([self.x, self.x], "stack", axis=0)
      self.assertEqual(stack_operation.splittable_dims, frozenset(["a", "b"]))
      self.assertEqual(stack_operation.unsplittable_dims, frozenset(["stack"]))
  
    def testUnstackOperation(self):
      unstack_operation = mtf.UnstackOperation(self.x, self.b_dim)
      self.assertEqual(unstack_operation.splittable_dims, frozenset(["a"]))
      self.assertEqual(unstack_operation.unsplittable_dims, frozenset(["b"]))
  
    def testEinsumOperation(self):
      x2 = mtf.zeros(self.mesh, mtf.Shape([self.a_dim, self.c_dim]))
      einsum_operation = mtf.EinsumOperation([self.x, x2],
                                             mtf.Shape([self.b_dim, self.c_dim]))
      self.assertEqual(einsum_operation.splittable_dims,
                       frozenset(["a", "b", "c"]))
      self.assertEqual(einsum_operation.unsplittable_dims, frozenset())
  
    def testConv2dOperations(self):
      conv_input = mtf.zeros(
          self.mesh,
          mtf.Shape([self.batch_dim, self.grid_h_dim, self.grid_w_dim,
                     self.in_dim]))
      conv_filter = mtf.zeros(
          self.mesh,
          mtf.Shape([self.filter_h_dim, self.filter_w_dim, self.in_dim,
                     self.out_dim]))
      strides = [1, 1, 1, 1]
      padding = "SAME"
  
      conv2d_operation = mtf.Conv2dOperation(conv_input, conv_filter, strides,
                                             padding)
      self.assertEqual(conv2d_operation.splittable_dims,
                       frozenset(["batch", "in", "out"]))
      self.assertEqual(conv2d_operation.unsplittable_dims,
                       frozenset(["filter_h", "filter_w", "grid_h", "grid_w"]))
  
      output = conv2d_operation.outputs[0]
      d_output = mtf.zeros(self.mesh, output.shape)
  
      conv2d_backprop_input_operation = mtf.Conv2or3dBackpropInputOperation(
          2, False, conv_input.shape, conv_filter, d_output, strides, padding)
      self.assertEqual(conv2d_backprop_input_operation.splittable_dims,
                       frozenset(["batch", "filter_h", "filter_w", "grid_h",
                                  "grid_w", "in", "out"]))
      self.assertEqual(conv2d_backprop_input_operation.unsplittable_dims,
                       frozenset())
  
      conv2d_backprop_filter_operation = mtf.Conv2or3dBackpropFilterOperation(
          2, False, conv_input, conv_filter.shape, d_output, strides, padding)
      self.assertEqual(conv2d_backprop_filter_operation.splittable_dims,
                       frozenset(["batch", "filter_h", "filter_w", "grid_h",
                                  "grid_w", "in", "out"]))
      self.assertEqual(conv2d_backprop_filter_operation.unsplittable_dims,
                       frozenset())
  
    def testShiftOperation(self):
      shift_operation = mtf.ShiftOperation(self.x, -5, self.b_dim, wrap=True)
      self.assertEqual(shift_operation.splittable_dims, frozenset(["a", "b"]))
      self.assertEqual(shift_operation.unsplittable_dims, frozenset())
  
    def testSliceOperation(self):
      slice_operation = mtf.SliceOperation(self.x, begin=3, size=4,
                                           slice_dim_name="b")
      self.assertEqual(slice_operation.splittable_dims, frozenset(["a"]))
      self.assertEqual(slice_operation.unsplittable_dims, frozenset(["b"]))
  
    def testPadOperation(self):
      pad_operation = mtf.PadOperation(self.x, [7, 2], "a")
      self.assertEqual(pad_operation.splittable_dims, frozenset(["b"]))
      self.assertEqual(pad_operation.unsplittable_dims, frozenset(["a"]))
  
    def testOneHotOperation(self):
      x = mtf.zeros(self.mesh, self.ab_shape, dtype=tf.int32)
      one_hot_operation = mtf.OneHotOperation(x, self.c_dim, 1, 0, dtype=tf.bool)
      self.assertEqual(one_hot_operation.splittable_dims,
                       frozenset(["a", "b", "c"]))
      self.assertEqual(one_hot_operation.unsplittable_dims, frozenset())
  
    def testImportOperation(self):
      tf_x = tf.zeros([5, 10])
      import_operation = mtf.ImportOperation(self.mesh, tf_x, self.ab_shape)
      self.assertEqual(import_operation.splittable_dims, frozenset(["a", "b"]))
      self.assertEqual(import_operation.unsplittable_dims, frozenset())
  
    def testImportLaidOutTensorOperation(self):
      laid_out_x = LaidOutTensor([self.x])
  
      import_laid_out_tensor_operation = mtf.ImportLaidOutTensorOperation(
          self.mesh, laid_out_x, self.ab_shape)
      self.assertEqual(import_laid_out_tensor_operation.splittable_dims,
                       frozenset())
      self.assertEqual(import_laid_out_tensor_operation.unsplittable_dims,
                       frozenset(["a", "b"]))
  
    def testVariableOperations(self):
      var = mtf.Variable(self.mesh,
                         "test_variable",
                         self.ab_shape,
                         mtf.VariableDType(tf.int32, tf.int32, tf.int32),
                         initializer=tf.zeros_initializer(),
                         trainable=True)
  
      self.assertEqual(var.splittable_dims, frozenset(["a", "b"]))
      self.assertEqual(var.unsplittable_dims, frozenset())
  
      read_variable = mtf.ReadVariable(var)
      self.assertEqual(read_variable.splittable_dims, frozenset(["a", "b"]))
      self.assertEqual(read_variable.unsplittable_dims, frozenset())
  
      assign = mtf.Assign([var], [self.x])
      self.assertEqual(assign.splittable_dims, frozenset(["a", "b"]))
      self.assertEqual(assign.unsplittable_dims, frozenset())
  
      depend = mtf.Depend(read_variable.outputs[0], [assign])
      self.assertEqual(depend.splittable_dims, frozenset(["a", "b"]))
      self.assertEqual(depend.unsplittable_dims, frozenset())
  
    def testConstant(self):
      constant = mtf.Constant(self.mesh, 0, self.ab_shape, dtype=tf.int32)
      self.assertEqual(constant.splittable_dims, frozenset(["a", "b"]))
      self.assertEqual(constant.unsplittable_dims, frozenset())
  
    def testStopGradient(self):
      stop_gradient = mtf.StopGradient(self.x)
      self.assertEqual(stop_gradient.splittable_dims, frozenset(["a", "b"]))
      self.assertEqual(stop_gradient.unsplittable_dims, frozenset())
  
    def testPrintOperation(self):
      print_operation = mtf.PrintOperation(self.x, [self.x], "Tensor x: ")
      self.assertEqual(print_operation.splittable_dims, frozenset(["a", "b"]))
      self.assertEqual(print_operation.unsplittable_dims, frozenset())
  
    def testReshapeOperation(self):
      reshape_operation = mtf.ReshapeOperation(
          self.x, mtf.Shape([mtf.Dimension("x", 25), mtf.Dimension("y", 2)]))
      self.assertEqual(reshape_operation.splittable_dims,
                       frozenset(["a", "b", "x", "y"]))
      self.assertEqual(reshape_operation.unsplittable_dims, frozenset())
  
    def testRandomOperation(self):
      random_operation = mtf.RandomOperation(self.mesh, self.ab_shape,
                                             tf.random_uniform)
      self.assertEqual(random_operation.splittable_dims, frozenset(["a", "b"]))
      self.assertEqual(random_operation.unsplittable_dims, frozenset())
  
    def testWhileLoopOperation(self):
      # This test case implements the following:
      # for i in range(10):
      #   x = x * 2
      i = mtf.constant(self.mesh, 0, mtf.Shape([]))
      cond_fn = lambda i, x: mtf.less(i, 10)
      body_fn = lambda i, x: [mtf.add(i, 1), mtf.multiply(x, 2)]
  
      while_loop_operation = mtf.WhileLoopOperation(cond_fn, body_fn, [i, self.x])
      self.assertEqual(while_loop_operation.splittable_dims,
                       frozenset(["a", "b"]))
      self.assertEqual(while_loop_operation.unsplittable_dims, frozenset())
  
  
  class NthSmallestTest(tf.test.TestCase):
  
    def testNthLargest(self):
      graph = mtf.Graph()
      mesh = mtf.Mesh(graph, "my_mesh")
      a_dim = mtf.Dimension("a", 6)
      b_dim = mtf.Dimension("b", 2)
      inputs = tf.constant([[1, 10],
                            [2, 9],
                            [3, 8],
                            [4, 7],
                            [5, 6],
                            [6, 5]])
      n = 1  # find second largest element (since n is zero-indexed)
      reduced_dim = a_dim
      expected_outputs = tf.constant([5, 9])
  
      mtf_inputs = mtf.import_tf_tensor(
          mesh, inputs, shape=mtf.Shape([a_dim, b_dim]))
      mtf_outputs = mtf.nth_largest_element(
          mtf_inputs, n, reduced_dim, "test_nth_largest")
      mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl(
          shape="all:2", layout="a:all", devices=["", ""])
      lowering = mtf.Lowering(graph, {mesh: mesh_impl})
      actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
      self.assertAllEqual(self.evaluate(actual_outputs),
                          self.evaluate(expected_outputs))
  
    def testNthSmallestReduceSecondDim(self):
      graph = mtf.Graph()
      mesh = mtf.Mesh(graph, "my_mesh")
      a_dim = mtf.Dimension("a", 6)
      b_dim = mtf.Dimension("b", 2)
      inputs = tf.constant([[1, 10],
                            [2, 9],
                            [3, 8],
                            [4, 7],
                            [5, 6],
                            [6, 5]])
      n = 0  # find smallest element (n is zero-indexed)
      reduced_dim = b_dim
      expected_outputs = tf.constant([1, 2, 3, 4, 5, 5])
  
      mtf_inputs = mtf.import_tf_tensor(
          mesh, inputs, shape=mtf.Shape([a_dim, b_dim]))
      mtf_outputs = mtf.nth_smallest_element(
          mtf_inputs, n, reduced_dim, "test_nth_smallest")
      mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl(
          shape="all:2", layout="a:all", devices=["", ""])
      lowering = mtf.Lowering(graph, {mesh: mesh_impl})
      actual_outputs = lowering.export_to_tf_tensor(mtf_outputs)
      self.assertAllEqual(self.evaluate(actual_outputs),
                          self.evaluate(expected_outputs))
  
  
  class TopKTest(tf.test.TestCase):
  
    def testTopK(self):
      graph = mtf.Graph()
      mesh = mtf.Mesh(graph, "my_mesh")
      a_dim = mtf.Dimension("a", 6)
      b_dim = mtf.Dimension("b", 2)
      inputs = tf.constant([[1, 10],
                            [2, 9],
                            [3, 8],
                            [4, 7],
                            [5, 6],
                            [6, 5]],
                           dtype=tf.float32)
      k_dim = mtf.Dimension("k", 2)
      d_values = tf.constant([[11, 12], [13, 14]], dtype=tf.float32)
      reduced_dim = a_dim
      expected_values = tf.constant([[6, 5], [10, 9]], dtype=tf.float32)
      expected_indices = tf.constant([[5, 4], [0, 1]])
      expected_d_inputs = tf.constant([[0, 13],
                                       [0, 14],
                                       [0, 0],
                                       [0, 0],
                                       [12, 0],
                                       [11, 0]],
                                      dtype=tf.float32)
  
      mtf_inputs = mtf.import_fully_replicated(
          mesh, inputs, shape=mtf.Shape([a_dim, b_dim]))
      mtf_d_values = mtf.import_tf_tensor(
          mesh, d_values, shape=mtf.Shape([b_dim, k_dim]))
      mtf_values, mtf_indices = mtf.top_k(mtf_inputs,
                                          reduced_dim=reduced_dim,
                                          k_dim=k_dim,
                                          name="test_nth_smallest")
      [mtf_d_inputs] = mtf.gradients([mtf_values], [mtf_inputs], [mtf_d_values])
      mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl(
          shape="rows:2,cols:2", layout="a:rows,b:cols", devices=["", "", "", ""])
      lowering = mtf.Lowering(graph, {mesh: mesh_impl})
      actual_values = lowering.export_to_tf_tensor(mtf_values)
      actual_indices = lowering.export_to_tf_tensor(mtf_indices)
      actual_d_inputs = lowering.export_to_tf_tensor(mtf_d_inputs)
      actual_inputs = lowering.export_to_tf_tensor(mtf_inputs)
      self.assertAllEqual(self.evaluate(actual_inputs),
                          self.evaluate(inputs))
      self.assertAllEqual(self.evaluate(actual_values),
                          self.evaluate(expected_values))
      self.assertAllEqual(self.evaluate(actual_indices),
                          self.evaluate(expected_indices))
      self.assertAllEqual(self.evaluate(actual_d_inputs),
                          self.evaluate(expected_d_inputs))
  
  
  if __name__ == "__main__":
    tf.disable_v2_behavior()
    tf.enable_eager_execution()
    tf.test.main()
  