ops_test.py

# coding=utf-8
# Copyright 2023 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
from mesh_tensorflow import test_utils as mtf_test_utils
import numpy as np
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]))

  @parameterized.parameters({
      "pool_fn": np.mean,
      "pool_fn_mtf": mtf.reduce_mean
  }, {
      "pool_fn": np.max,
      "pool_fn_mtf": mtf.reduce_max
  }, {
      "pool_fn": np.min,
      "pool_fn_mtf": mtf.reduce_min
  })
  def testPoolTensor1d(self, pool_fn, pool_fn_mtf):
    converter = mtf_test_utils.NumpyConverter()
    pool_size = 2
    x = np.random.randn(2, 3, 4, 5)
    expected = np.empty(shape=[2, 3, 2, 5])
    expected[:, :, 0, :] = pool_fn(x[:, :, 0:2, :], axis=2)
    expected[:, :, 1, :] = pool_fn(x[:, :, 2:4, :], axis=2)

    x_mtf = converter.convert_np_array_to_mtf_tensor(x, dtype=tf.float32)
    pooled_mtf = mtf.pool_tensor_1d(
        x_mtf,
        pool_dim=x_mtf.shape.dims[2],
        reduce_fn=pool_fn_mtf,
        pool_size=pool_size)
    actual = converter.convert_mtf_tensor_to_np_array(pooled_mtf)
    self.assertAllClose(expected, actual)

  @parameterized.parameters({"pool_size": 2}, {"pool_size": 3})
  def testStrideTensor1d(self, pool_size):
    converter = mtf_test_utils.NumpyConverter()
    x = np.random.randint(0, 100, size=[2, 3, 6, 5])
    x_mtf = converter.convert_np_array_to_mtf_tensor(x)
    expected = x[:, :, range(0, x.shape[2], pool_size), :]
    strided_mtf = mtf.stride_tensor_1d(
        x_mtf, pool_dim=x_mtf.shape.dims[2], pool_size=pool_size)
    actual = converter.convert_mtf_tensor_to_np_array(strided_mtf)
    self.assertAllEqual(expected, actual)

  def testReduceFirst(self):
    converter = mtf_test_utils.NumpyConverter()
    x = np.random.randint(0, 100, size=[2, 3, 6, 5])
    x_mtf = converter.convert_np_array_to_mtf_tensor(x)
    expected = x[:, :, 0, :]
    reduced_mtf = mtf.reduce_first(x_mtf, reduced_dim=x_mtf.shape.dims[2])
    actual = converter.convert_mtf_tensor_to_np_array(reduced_mtf)
    self.assertAllEqual(expected, actual)


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))


class RecomputeGradTest(tf.test.TestCase):

  def testRecomputeGrad(self):
    graph = mtf.Graph()
    mesh = mtf.Mesh(graph, "my_mesh")
    # let's differentiate x^2 + x
    # dy/dx = 2x+1
    def x_squared_plus_x(x):
      return x * x + x
    x = tf.constant([5, 10], dtype=tf.float32)
    dy = tf.constant([2, 3], dtype=tf.float32)
    two = mtf.Dimension("two", 2)
    expected_y = tf.constant([30, 110], dtype=tf.float32)
    expected_dx = tf.constant([22, 63], dtype=tf.float32)
    mtf_x = mtf.import_fully_replicated(
        mesh, x, shape=mtf.Shape([two]))
    mtf_dy = mtf.import_tf_tensor(
        mesh, dy, shape=mtf.Shape([two]))
    mtf_y = mtf.recompute_grad(x_squared_plus_x, [mtf_x])
    [mtf_dx] = mtf.gradients([mtf_y], [mtf_x], [mtf_dy])
    mesh_impl = mtf.placement_mesh_impl.PlacementMeshImpl(
        shape="processors:2", layout="two:processors", devices=["", ""])
    lowering = mtf.Lowering(graph, {mesh: mesh_impl})
    actual_y = lowering.export_to_tf_tensor(mtf_y)
    actual_dx = lowering.export_to_tf_tensor(mtf_dx)
    self.assertAllEqual(self.evaluate(actual_y),
                        self.evaluate(expected_y))
    self.assertAllEqual(self.evaluate(actual_dx),
                        self.evaluate(expected_dx))


if __name__ == "__main__":
  tf.disable_v2_behavior()
  tf.enable_eager_execution()
  tf.test.main()