layout_optimizer_test.py

# Copyright 2017 The TensorFlow Authors. All Rights Reserved.
#
# 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 Grappler LayoutOptimizer."""

import numpy as np

from tensorflow.core.protobuf import config_pb2
from tensorflow.core.protobuf import device_properties_pb2
from tensorflow.core.protobuf import rewriter_config_pb2
from tensorflow.core.protobuf import saver_pb2
from tensorflow.python.client import session
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import random_seed
from tensorflow.python.framework import test_util
from tensorflow.python.grappler import cluster as gcluster
from tensorflow.python.grappler import tf_optimizer
from tensorflow.python.layers import convolutional as conv_layers
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import gen_array_ops
from tensorflow.python.ops import gen_math_ops
from tensorflow.python.ops import gen_nn_ops
from tensorflow.python.ops import map_fn
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import nn
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import variables
from tensorflow.python.platform import test
from tensorflow.python.training import gradient_descent
from tensorflow.python.training import saver as saver_lib


def _weight(shape):
  """Generates a weight of a given shape."""
  return random_ops.truncated_normal(shape, seed=0, stddev=0.1)


def _bias(shape):
  """Generates a bias of a given shape."""
  return constant_op.constant(0.1, shape=shape)


def _conv2d(x, w):
  """Returns a 2d convolution layer with full stride."""
  return nn.conv2d(x, w, strides=[1, 1, 1, 1], padding='SAME')


def _max_pool_2x2(x):
  """Downsamples a feature map by 2X."""
  return nn.max_pool(
      x, ksize=[1, 2, 2, 1], strides=[1, 2, 2, 1], padding='SAME')


# Taken from tensorflow/examples/tutorials/mnist/mnist_deep.py
def _two_layer_model(x):
  x_image = array_ops.reshape(x, [-1, 28, 28, 1])
  w_conv1 = _weight([5, 5, 1, 32])
  b_conv1 = _bias([32])
  h_conv1 = nn.relu(_conv2d(x_image, w_conv1) + b_conv1)
  h_pool1 = _max_pool_2x2(h_conv1)
  w_conv2 = _weight([5, 5, 32, 64])
  b_conv2 = _bias([64])
  h_conv2 = nn.relu(_conv2d(h_pool1, w_conv2) + b_conv2)
  h_pool2 = _max_pool_2x2(h_conv2)
  return h_pool2


def _model_with_second_port():
  random_seed.set_random_seed(0)
  x = random_ops.truncated_normal([2, 5, 5, 4], seed=0)
  scale = constant_op.constant(0.1, shape=[4])
  offset = constant_op.constant(0.3, shape=[4])
  y, mean, _ = nn.fused_batch_norm(x, scale, offset)
  mul = math_ops.add(y, mean)
  output = array_ops.identity(mul)
  return output


def _model_with_branch(x):
  x_image = array_ops.reshape(x, [-1, 28, 28, 1])
  w_conv1 = _weight([5, 5, 1, 32])
  w_conv2 = _weight([5, 5, 1, 32])
  c_conv1 = _conv2d(x_image, w_conv1)
  c_conv2 = _conv2d(x_image, w_conv2)
  add = math_ops.add(c_conv1, c_conv2)
  return add


def _model_with_vec_and_4d(x):
  x_image = array_ops.reshape(x, [-1, 28, 28, 1])
  w_conv1 = _weight([5, 5, 1, 32])
  c_conv1 = _conv2d(x_image, w_conv1)
  vector = constant_op.constant(6.4, shape=[32])
  add = math_ops.add(c_conv1, vector)
  return add


def _loop():
  random_seed.set_random_seed(0)
  x1 = random_ops.truncated_normal([1, 784], seed=0)
  x2 = random_ops.truncated_normal([1, 784], seed=0)
  x3 = random_ops.truncated_normal([1, 784], seed=0)
  x4 = random_ops.truncated_normal([1, 784], seed=0)
  elems = (x1, x2, x3, x4)
  outputs = map_fn.map_fn(_two_layer_model, elems, dtype=dtypes.float32)
  return outputs


def _loop_with_branch():
  random_seed.set_random_seed(0)
  x1 = random_ops.truncated_normal([1, 784], seed=0)
  x2 = random_ops.truncated_normal([1, 784], seed=0)
  x3 = random_ops.truncated_normal([1, 784], seed=0)
  x4 = random_ops.truncated_normal([1, 784], seed=0)
  elems = (x1, x2, x3, x4)
  outputs = map_fn.map_fn(_model_with_branch, elems, dtype=dtypes.float32)
  return outputs


def _loop_with_vec_and_4d():
  random_seed.set_random_seed(0)
  x1 = random_ops.truncated_normal([1, 784], seed=0)
  x2 = random_ops.truncated_normal([1, 784], seed=0)
  x3 = random_ops.truncated_normal([1, 784], seed=0)
  x4 = random_ops.truncated_normal([1, 784], seed=0)
  elems = (x1, x2, x3, x4)
  outputs = map_fn.map_fn(_model_with_vec_and_4d, elems, dtype=dtypes.float32)
  return outputs


def _get_config(layout_optimizer=True):
  if layout_optimizer:
    rewrite_options = rewriter_config_pb2.RewriterConfig(
        layout_optimizer=rewriter_config_pb2.RewriterConfig.ON,
        # do not remove duplicated nodes
        arithmetic_optimization=rewriter_config_pb2.RewriterConfig.OFF)
  else:
    rewrite_options = rewriter_config_pb2.RewriterConfig(
        layout_optimizer=rewriter_config_pb2.RewriterConfig.OFF,
        # do not remove duplicated nodes
        arithmetic_optimization=rewriter_config_pb2.RewriterConfig.OFF)
  rewrite_options.min_graph_nodes = -1
  graph_options = config_pb2.GraphOptions(
      rewrite_options=rewrite_options, build_cost_model=1)
  config = config_pb2.ConfigProto(graph_options=graph_options)
  config.graph_options.optimizer_options.opt_level = -1
  return config


def _simple_metagraph(depthwise=False):
  random_seed.set_random_seed(0)
  x = variables.Variable(random_ops.truncated_normal([1, 200, 200, 3], seed=0))
  conv = conv_layers.separable_conv2d if depthwise else conv_layers.conv2d
  y = conv(x, 32, [3, 3])
  z = conv(y, 32, [3, 3])
  optimizer = gradient_descent.GradientDescentOptimizer(1e-4)
  loss = math_ops.reduce_mean(z)
  train_op = optimizer.minimize(loss)
  graph = ops.get_default_graph()
  graph.add_to_collection('train_op', train_op)
  meta_graph = saver_lib.export_meta_graph(graph_def=graph.as_graph_def())
  return meta_graph


def _get_cluster():
  named_device = device_properties_pb2.NamedDevice()
  named_device.name = '/GPU:0'
  named_device.properties.type = 'GPU'
  named_device.properties.num_cores = 24
  named_device.properties.frequency = 1000
  named_device.properties.environment['architecture'] = '4'
  cluster = gcluster.Cluster(devices=[named_device])
  return cluster


def _is_transpose(node):
  return node.endswith('TransposeNHWCToNCHW-LayoutOptimizer') or node.endswith(
      'TransposeNCHWToNHWC-LayoutOptimizer') or node.endswith(
          'TransposeNDHWCToNCDHW-LayoutOptimizer') or node.endswith(
              'TransposeNCDHWToNDHWC-LayoutOptimizer')


def _is_permute(node):
  return node.endswith('VecPermuteNHWCToNCHW-LayoutOptimizer') or node.endswith(
      'VecPermuteNCHWToNHWC-LayoutOptimizer')


@test_util.for_all_test_methods(test_util.no_xla_auto_jit,
                                'Test does not apply in XLA setting')
class LayoutOptimizerTest(test.TestCase):
  """Tests the Grappler layout optimizer."""

  def _assert_trans_nchw_to_nhwc(self, name, nodes):
    self.assertIn(name + '-TransposeNCHWToNHWC-LayoutOptimizer', nodes)

  def _assert_trans_nhwc_to_nchw(self, name, nodes):
    self.assertIn(name + '-TransposeNHWCToNCHW-LayoutOptimizer', nodes)

  def _assert_trans_ncdhw_to_ndhwc(self, name, nodes):
    self.assertIn(name + '-TransposeNCDHWToNDHWC-LayoutOptimizer', nodes)

  def _assert_trans_ndhwc_to_ncdhw(self, name, nodes):
    self.assertIn(name + '-TransposeNDHWCToNCDHW-LayoutOptimizer', nodes)

  def _assert_map_nhwc_to_nchw(self, name, nodes):
    self.assertIn(name + '-DimMapNHWCToNCHW-LayoutOptimizer', nodes)

  def _assert_map_ndhwc_to_ncdhw(self, name, nodes):
    self.assertIn(name + '-DataFormatDimMapNDHWCToNCDHW-LayoutOptimizer', nodes)

  def _assert_vec_nchw_to_nhwc(self, name, nodes):
    self.assertIn(name + '-VecPermuteNCHWToNHWC-LayoutOptimizer', nodes)

  def _assert_vec_nhwc_to_nchw(self, name, nodes):
    self.assertIn(name + '-VecPermuteNHWCToNCHW-LayoutOptimizer', nodes)

  def _assert_vec_ncdhw_to_ndhwc(self, name, nodes):
    self.assertIn(name + '-DataFormatVecPermuteNCDHWToNDHWC-LayoutOptimizer',
                  nodes)

  def _assert_vec_ndhwc_to_ncdhw(self, name, nodes):
    self.assertIn(name + '-DataFormatVecPermuteNDHWCToNCDHW-LayoutOptimizer',
                  nodes)

  def _train(self, checkpoint_path, layout_optimizer=False, restore=False):
    ops.reset_default_graph()
    graph = ops.get_default_graph()
    with session.Session(
        config=_get_config(layout_optimizer), graph=graph) as sess:
      batch = 2
      height = 6
      width = 7
      input_channels = 3
      shape = [batch, height, width, input_channels]
      image = array_ops.placeholder(dtype='float32', shape=shape)
      conv1 = conv_layers.conv2d(image, 32, [3, 3])
      conv2 = conv_layers.conv2d(conv1, 32, [3, 3])
      optimizer = gradient_descent.GradientDescentOptimizer(0.01)
      loss = math_ops.reduce_mean(conv2)
      train_op = optimizer.minimize(loss)
      saver = saver_lib.Saver(write_version=saver_pb2.SaverDef.V2)

      if restore:
        saver.restore(sess, checkpoint_path)
      else:
        self.evaluate(variables.global_variables_initializer())

      np.random.seed(0)
      for _ in range(2):
        image_val = np.random.rand(*shape).astype(np.float32)
        sess.run([loss, train_op], feed_dict={image: image_val})

      if restore:
        all_vars = ops.get_collection(ops.GraphKeys.GLOBAL_VARIABLES)
        all_vars_values = [var.eval(session=sess) for var in all_vars]
        return all_vars_values
      else:
        saver.save(sess, checkpoint_path)

  @test_util.deprecated_graph_mode_only
  def testTwoConvLayers(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      output = _two_layer_model(x)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = self.evaluate(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Four transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('Relu_1-0-0', nodes)

      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testSplitWithNonConstAxis(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      dim = array_ops.placeholder(dtype='int32')
      split = array_ops.split(conv, 2, axis=dim)
      scale = constant_op.constant(0.1, shape=[32])
      offset = constant_op.constant(0.3, shape=[32])
      bn0 = nn.fused_batch_norm(split[0], scale, offset)
      bn1 = nn.fused_batch_norm(split[1], scale, offset)
      add = bn0[0] + bn1[0]
      output = array_ops.identity(add)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = sess.run(output, feed_dict={dim: 3})

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata, feed_dict={dim: 3})

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('add_2-0-0', nodes)
      self._assert_map_nhwc_to_nchw('split-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testSplitVWithNonConstAxis(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      dim = array_ops.placeholder(dtype='int32')
      sizes = constant_op.constant([50, 10, 4], shape=[3])
      split = gen_array_ops.split_v(
          value=conv, size_splits=sizes, axis=dim, num_split=3)
      output = math_ops.reduce_sum(split[0])

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = sess.run(output, feed_dict={dim: 3})

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata, feed_dict={dim: 3})

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Four transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('SplitV-0-0', nodes)
      self._assert_map_nhwc_to_nchw('SplitV-2', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testPadWithConstPaddings(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      paddings_val = [[1, 2], [3, 4], [5, 6], [7, 8]]
      paddings = constant_op.constant(
          paddings_val, dtype='int32', name='PaddingsConst')
      pad = array_ops.pad(conv, paddings)
      output = array_ops.identity(pad)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = self.evaluate(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Four transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('Pad-0-0', nodes)
      self.assertIn('Pad-1-LayoutOptimizer', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testReduceSum(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      reduce_sum = math_ops.reduce_sum(conv)
      output = array_ops.identity(reduce_sum)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = self.evaluate(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Three transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 1
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testCast(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      cast = math_ops.cast(conv, dtype='bool')
      output = array_ops.identity(cast)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = self.evaluate(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Four transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('Cast-0-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testSqueeze(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      reduce_sum = math_ops.reduce_sum(conv, axis=[1, 2])
      squeeze = array_ops.squeeze(reduce_sum)
      output = array_ops.identity(squeeze)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = self.evaluate(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Three transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 1
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testSqueezeAlongHW(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      reduce_sum = math_ops.reduce_sum(conv, axis=[1, 2], keepdims=True)
      squeeze = array_ops.squeeze(reduce_sum, axis=[1, 2])
      output = array_ops.identity(squeeze)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = self.evaluate(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Three transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 1
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testSqueezeAlongNHW(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      reduce_sum = math_ops.reduce_sum(conv, axis=[0, 1, 2], keepdims=True)
      squeeze = array_ops.squeeze(reduce_sum, axis=[0, 1, 2])
      output = array_ops.identity(squeeze)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = self.evaluate(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Three transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 1
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testReduceSumAlongHWC(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      reduce_sum = math_ops.reduce_sum(conv, axis=[1, 2, 3])
      output = array_ops.identity(reduce_sum)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = self.evaluate(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Three transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 1
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testReduceSumAlongNHW(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      reduce_sum = math_ops.reduce_sum(conv, axis=[0, 1, 2])
      output = array_ops.identity(reduce_sum)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = self.evaluate(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Three transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 1
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testReduceSumAlongC(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      reduce_sum = math_ops.reduce_sum(conv, axis=[3])
      output = array_ops.identity(reduce_sum)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = self.evaluate(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Three transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 1
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testReduceSumAlongCKeepDims(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      reduce_sum = math_ops.reduce_sum(conv, axis=[3], keepdims=True)
      output = array_ops.identity(reduce_sum)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = self.evaluate(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Four transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('Sum-0-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testReduceSumAlongHKeepDims(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      reduce_sum = math_ops.reduce_sum(conv, axis=[2], keepdims=True)
      output = array_ops.identity(reduce_sum)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = self.evaluate(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Four transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testReduceSumAlongWCKeepDims(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      reduce_sum = math_ops.reduce_sum(conv, axis=[2, 3], keepdims=True)
      output = array_ops.identity(reduce_sum)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = self.evaluate(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Four transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testConcatWithControlDependency(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      axis = constant_op.constant(3)
      var = variables.Variable(3)
      assign = state_ops.assign(var, 6)
      with ops.control_dependencies([assign]):
        concat = array_ops.concat([conv, conv], axis)
      output = array_ops.identity(concat)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = self.evaluate(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Four transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('concat-0-0', nodes)
      self.assertIn('concat-2-LayoutOptimizer', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testConcatWithControlDependencyFor5DTensor(self):
    if not test.is_gpu_available(cuda_only=True):
      self.skipTest('GPU required')
    random_seed.set_random_seed(0)
    x = random_ops.truncated_normal([2, 2, 14, 14, 1], seed=0)
    w = random_ops.truncated_normal([2, 2, 2, 1, 2], seed=0)
    strides = [1, 1, 1, 1, 1]
    y = gen_nn_ops.conv3d(x, w, strides, 'SAME')
    axis = constant_op.constant(4)
    var = variables.Variable(3)
    assign = state_ops.assign(var, 6)
    with ops.control_dependencies([assign]):
      concat = array_ops.concat([y, y], axis)
    output = array_ops.identity(concat)

    with session.Session(config=_get_config(False)) as sess:
      output_val_ref = self.evaluate(output)

    with session.Session(config=_get_config()) as sess:
      metadata = config_pb2.RunMetadata()
      output_val = sess.run(output, run_metadata=metadata)

    nodes = []
    num_transposes = 0
    for node in metadata.cost_graph.node:
      if _is_transpose(node.name):
        num_transposes += 1
      nodes.append(node.name)

    # Four transposes were initially added in the Expand phase of
    # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
    expected_num_transposes = 2
    self.assertEqual(expected_num_transposes, num_transposes)
    self._assert_trans_ndhwc_to_ncdhw('Conv3D-0', nodes)
    self._assert_trans_ncdhw_to_ndhwc('concat-0-0', nodes)
    self._assert_map_ndhwc_to_ncdhw('concat-2', nodes)
    self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testFill(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = array_ops.placeholder(dtype='float32')
      conv = _two_layer_model(x)
      shape = array_ops.shape(conv)
      scalar = array_ops.constant(5.7)
      fill = array_ops.fill(shape, scalar)
      output = array_ops.identity(fill)

      x_val = [3.4] * 784
      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = sess.run(output, feed_dict={x: x_val})

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(
            output, run_metadata=metadata, feed_dict={
                x: x_val
            })

      nodes = []
      num_transposes = 0
      num_vec_permute = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        if _is_permute(node.name):
          num_vec_permute += 1
        nodes.append(node.name)

      # Four transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      # Two vector permute nodes were initially added in the Expand phase of
      # LayoutOptimizer; they cancelled out each other in the Collapse phase.
      expected_vec_permute = 0
      self.assertEqual(expected_vec_permute, num_vec_permute)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('Fill-0-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testTile(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      multiple = array_ops.placeholder(dtype='int32')
      tile = array_ops.tile(conv, multiple)
      output = array_ops.identity(tile)

      multiple_val = [2, 3, 4, 1]
      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = sess.run(output, feed_dict={multiple: multiple_val})

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(
            output, run_metadata=metadata, feed_dict={
                multiple: multiple_val
            })

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Four transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('Tile-0-0', nodes)
      self._assert_vec_nhwc_to_nchw('Tile-1', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testReverseWithConstDims(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      dims = constant_op.constant([3, 1], name='DimsConst')
      reverse = array_ops.reverse(conv, dims)
      output = array_ops.identity(reverse)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = self.evaluate(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Four transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('ReverseV2-0-0', nodes)
      self.assertIn('ReverseV2-1-LayoutOptimizer', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testReverseWithNonConstDims(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      dims = array_ops.placeholder(dtype='int32')
      reverse = array_ops.reverse(conv, dims)
      output = array_ops.identity(reverse)

      dims_val = [2, 3]
      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = sess.run(output, feed_dict={dims: dims_val})

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(
            output, run_metadata=metadata, feed_dict={
                dims: dims_val
            })

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Four transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('ReverseV2-0-0', nodes)
      self._assert_map_nhwc_to_nchw('ReverseV2-1', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testSelectOp(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      add = math_ops.add(conv, conv)
      mean = math_ops.reduce_mean(conv)
      condition = math_ops.less(conv, mean)
      select = gen_math_ops.select(condition, conv, add)
      output = array_ops.identity(select)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = self.evaluate(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('Select-0-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testSelectOpConditionUnknownShape(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      add = math_ops.add(conv, conv)
      condition = array_ops.placeholder(dtype='bool')
      select = gen_math_ops.select(condition, conv, add)
      output = array_ops.identity(select)

      condition_val = np.zeros((1, 7, 7, 64))
      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = sess.run(output, feed_dict={condition: condition_val})

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(
            output, run_metadata=metadata, feed_dict={condition: condition_val})

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      expected_num_transposes = 3
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testSelectOpScalarCondition(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      add = math_ops.add(conv, conv)
      condition = constant_op.constant(True)
      select = gen_math_ops.select(condition, conv, add)
      output = array_ops.identity(select)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = self.evaluate(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('Select-0-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testPadWithNonConstPaddings(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      paddings = array_ops.placeholder(dtype='int32')
      pad = array_ops.pad(conv, paddings)
      output = array_ops.identity(pad)

      paddings_val = [[1, 2], [3, 4], [5, 6], [7, 8]]
      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = sess.run(output, feed_dict={paddings: paddings_val})

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(
            output, run_metadata=metadata, feed_dict={
                paddings: paddings_val
            })

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Four transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('Pad-0-0', nodes)
      self._assert_vec_nhwc_to_nchw('Pad-1', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testMaxPoolV2(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      ksize = constant_op.constant([1, 2, 3, 1], shape=[4])
      strides = array_ops.placeholder(dtype='int32', shape=[4])
      max_pool = gen_nn_ops.max_pool_v2(conv, ksize, strides, 'VALID')
      output = array_ops.identity(max_pool)

      strides_val = [1, 3, 2, 1]
      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = sess.run(output, feed_dict={strides: strides_val})

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(
            output, run_metadata=metadata, feed_dict={
                strides: strides_val
            })

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('MaxPoolV2-0-0', nodes)
      self._assert_vec_nhwc_to_nchw('MaxPoolV2-2', nodes)
      self.assertIn('MaxPoolV2-1-LayoutOptimizer', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testMaxPoolGradV2(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      ksize = constant_op.constant([1, 2, 3, 1], shape=[4])
      strides = array_ops.placeholder(dtype='int32', shape=[4])
      max_pool_grad = gen_nn_ops.max_pool_grad_v2(conv, conv, conv, ksize,
                                                  strides, 'VALID')
      output = array_ops.identity(max_pool_grad)

      strides_val = [1, 3, 2, 1]
      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = sess.run(output, feed_dict={strides: strides_val})

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(
            output, run_metadata=metadata, feed_dict={
                strides: strides_val
            })

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('MaxPoolGradV2-0-0', nodes)
      self._assert_vec_nhwc_to_nchw('MaxPoolGradV2-4', nodes)
      self.assertIn('MaxPoolGradV2-3-LayoutOptimizer', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testLeakyRelu(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([4, 14, 14, 1], seed=0)
      w = random_ops.truncated_normal([2, 2, 1, 2], seed=0)
      y = nn.conv2d(x, w, strides=[1, 1, 1, 1], padding='SAME')
      y = nn.leaky_relu(y, alpha=0.2)
      output = array_ops.identity(y)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = sess.run(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nchw_to_nhwc('LeakyRelu-0-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testLeakyReluGrad(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([4, 14, 14, 1], seed=0)
      w = random_ops.truncated_normal([2, 2, 1, 1], seed=0)
      y = nn.conv2d(x, w, strides=[1, 1, 1, 1], padding='SAME')
      y = gen_nn_ops.leaky_relu_grad(y, x, alpha=0.2)
      output = array_ops.identity(y)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = sess.run(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      expected_num_transposes = 3
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('LeakyReluGrad-1', nodes)
      self._assert_trans_nchw_to_nhwc('LeakyReluGrad-0-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testLeakyReluGradFor5DTensors(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 4, 2, 3, 3], seed=0)
      w = random_ops.truncated_normal([2, 2, 2, 3, 3], seed=0)
      y = gen_nn_ops.conv3d(x, w, [1, 1, 1, 1, 1], 'SAME')
      y = gen_nn_ops.leaky_relu_grad(y, x, alpha=0.2)
      output = array_ops.identity(y)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = sess.run(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      expected_num_transposes = 3
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_ndhwc_to_ncdhw('LeakyReluGrad-1', nodes)
      self._assert_trans_ncdhw_to_ndhwc('LeakyReluGrad-0-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testReduceOpsFor5DTensors(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 4, 2, 3, 3], seed=0)
      w = random_ops.truncated_normal([2, 2, 2, 3, 3], seed=0)
      conv3d = gen_nn_ops.conv3d(x, w, [1, 1, 1, 1, 1], 'SAME')
      y = math_ops.reduce_mean(conv3d, [0, 1, 2, 3], keepdims=True)
      output = array_ops.identity(y)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = sess.run(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # The reduce op Mean needs to dim map the input reduce index to NCDHW.
      # Then, the output needs to be tranposed back to NDHWC.
      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_ndhwc_to_ncdhw('Conv3D-0', nodes)
      self._assert_map_ndhwc_to_ncdhw('Mean-1', nodes)
      self._assert_trans_ncdhw_to_ndhwc('Mean-0-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testBinaryOpsFor5DTensors(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 4, 2, 3, 3], seed=0)
      w = random_ops.truncated_normal([2, 2, 2, 3, 3], seed=0)
      mean = random_ops.truncated_normal([1, 1, 1, 1, 3], seed=0)
      variance = random_ops.truncated_normal([1, 1, 1, 1, 3], seed=0)
      gamma = random_ops.truncated_normal([1, 1, 1, 1, 3], seed=0)
      beta = random_ops.truncated_normal([1, 1, 1, 1, 3], seed=0)
      conv3d = gen_nn_ops.conv3d(x, w, [1, 1, 1, 1, 1], 'SAME')
      y = nn.batch_normalization(
          conv3d,
          mean=mean,
          variance=variance,
          scale=gamma,
          offset=beta,
          variance_epsilon=0.001)
      output = array_ops.identity(y)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = sess.run(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # The binary ops mul_1 and add_1 in batch norm need to transpose one of
      # the two inputs to NCDHW. The other input has already been tranposed via
      # Conv3D.
      expected_num_transposes = 4
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_ndhwc_to_ncdhw('Conv3D-0', nodes)
      self._assert_trans_ndhwc_to_ncdhw('batchnorm/mul_1-1', nodes)
      self._assert_trans_ndhwc_to_ncdhw('batchnorm/add_1-1', nodes)
      self._assert_trans_ncdhw_to_ndhwc('batchnorm/add_1-0-0', nodes)

  @test_util.deprecated_graph_mode_only
  def testBatchNorm3D(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x_3d = random_ops.truncated_normal([1, 4, 2, 3, 3], seed=0)
      filters = random_ops.truncated_normal([2, 2, 2, 3, 3], seed=0)
      strides_val = [1, 1, 1, 1, 1]
      scale = constant_op.constant(0.1, shape=[3])
      offset = constant_op.constant(0.3, shape=[3])
      conv3d = gen_nn_ops.conv3d(x_3d, filters, strides_val, 'SAME')
      y, _, _ = nn.fused_batch_norm(conv3d, scale, offset, data_format='NDHWC')
      output = array_ops.identity(y)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = sess.run(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_ndhwc_to_ncdhw('Conv3D-0', nodes)
      self._assert_trans_ncdhw_to_ndhwc('FusedBatchNormV3-0-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testBatchNormGrad3D(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x_3d = random_ops.truncated_normal([1, 4, 2, 3, 3], seed=0)
      filters = random_ops.truncated_normal([2, 2, 2, 3, 3], seed=0)
      strides_val = [1, 1, 1, 1, 1]
      scale = constant_op.constant(0.1, shape=[3])
      offset = constant_op.constant(0.3, shape=[3])
      mean = constant_op.constant(0.1, shape=[3])
      variance = constant_op.constant(0.3, shape=[3])
      conv3d = gen_nn_ops.conv3d(x_3d, filters, strides_val, 'SAME')
      y, running_mean, running_var, r0, r1, r2 = gen_nn_ops.fused_batch_norm_v3(
          conv3d,
          scale,
          offset,
          mean,
          variance,
          epsilon=1.001e-5,
          exponential_avg_factor=1.0,
          data_format='NDHWC',
          is_training=True,
          name='batch_norm')
      dx, dscale, doffset, _, _ = gen_nn_ops.fused_batch_norm_grad_v3(
          y,
          x_3d,
          scale,
          r0,
          r1,
          r2,
          epsilon=1.001e-5,
          data_format='NDHWC',
          is_training=True)
      output = array_ops.identity(dx)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = sess.run(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      expected_num_transposes = 3
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_ndhwc_to_ncdhw('Conv3D-0', nodes)
      self._assert_trans_ndhwc_to_ncdhw('FusedBatchNormGradV3-1', nodes)
      self._assert_trans_ncdhw_to_ndhwc('FusedBatchNormGradV3-0-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testConv3D(self):
    if not test.is_gpu_available(cuda_only=True):
      self.skipTest('GPU required')
    random_seed.set_random_seed(0)
    x = random_ops.truncated_normal([2, 2, 14, 14, 1], seed=0)
    w = random_ops.truncated_normal([2, 2, 2, 1, 2], seed=0)
    strides = [1, 1, 1, 1, 1]
    y = gen_nn_ops.conv3d(x, w, strides, 'SAME')
    output = array_ops.identity(y)

    with session.Session(config=_get_config(False)) as sess:
      output_val_ref = sess.run(output)

    with session.Session(config=_get_config()) as sess:
      metadata = config_pb2.RunMetadata()
      output_val = sess.run(output, run_metadata=metadata)

    nodes = []
    num_transposes = 0
    for node in metadata.cost_graph.node:
      if _is_transpose(node.name):
        num_transposes += 1
      nodes.append(node.name)

    expected_num_transposes = 2
    self.assertEqual(expected_num_transposes, num_transposes)
    self._assert_trans_ndhwc_to_ncdhw('Conv3D-0', nodes)
    self._assert_trans_ncdhw_to_ndhwc('Conv3D-0-0', nodes)
    self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testConv3DBackpropInput(self):
    if not test.is_gpu_available(cuda_only=True):
      self.skipTest('GPU required')
    random_seed.set_random_seed(0)
    dy = random_ops.truncated_normal([2, 2, 14, 14, 1], seed=0)
    w = random_ops.truncated_normal([2, 2, 2, 1, 1], seed=0)
    strides = [1, 1, 1, 1, 1]
    x_shape = array_ops.shape(dy)
    dx = gen_nn_ops.conv3d_backprop_input_v2(x_shape, w, dy, strides, 'SAME')
    output = array_ops.identity(dx)

    with session.Session(config=_get_config(False)) as sess:
      output_val_ref = sess.run(output)

    with session.Session(config=_get_config()) as sess:
      metadata = config_pb2.RunMetadata()
      output_val = sess.run(output, run_metadata=metadata)

    nodes = []
    num_transposes = 0
    for node in metadata.cost_graph.node:
      if _is_transpose(node.name):
        num_transposes += 1
      nodes.append(node.name)

    expected_num_transposes = 2
    self.assertEqual(expected_num_transposes, num_transposes)
    self._assert_vec_ndhwc_to_ncdhw('Conv3DBackpropInputV2-0', nodes)
    self._assert_trans_ndhwc_to_ncdhw('Conv3DBackpropInputV2-2', nodes)
    self._assert_trans_ncdhw_to_ndhwc('Conv3DBackpropInputV2-0-0', nodes)
    self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testConv3DBackpropFilter(self):
    if not test.is_gpu_available(cuda_only=True):
      self.skipTest('GPU required')
    random_seed.set_random_seed(0)
    x = random_ops.truncated_normal([2, 2, 14, 14, 1], seed=0)
    dy = random_ops.truncated_normal([2, 2, 14, 14, 1], seed=0)
    strides = [1, 1, 1, 1, 1]
    w_shape = constant_op.constant([2, 2, 2, 1, 1], shape=[5])
    dw = gen_nn_ops.conv3d_backprop_filter_v2(x, w_shape, dy, strides, 'SAME')
    output = array_ops.identity(dw)

    with session.Session(config=_get_config(False)) as sess:
      output_val_ref = sess.run(output)

    with session.Session(config=_get_config()) as sess:
      metadata = config_pb2.RunMetadata()
      output_val = sess.run(output, run_metadata=metadata)

    nodes = []
    num_transposes = 0
    for node in metadata.cost_graph.node:
      if _is_transpose(node.name):
        num_transposes += 1
      nodes.append(node.name)

    expected_num_transposes = 2
    self.assertEqual(expected_num_transposes, num_transposes)
    self._assert_trans_ndhwc_to_ncdhw('Conv3DBackpropFilterV2-0', nodes)
    self._assert_trans_ndhwc_to_ncdhw('Conv3DBackpropFilterV2-2', nodes)
    self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testBiasAddFor5DTensor(self):
    if not test.is_gpu_available(cuda_only=True):
      self.skipTest('GPU required')
    random_seed.set_random_seed(0)
    x = random_ops.truncated_normal([2, 2, 14, 14, 1], seed=0)
    w = random_ops.truncated_normal([2, 2, 2, 1, 2], seed=0)
    b = random_ops.truncated_normal([2], seed=0)
    strides = [1, 1, 1, 1, 1]
    y = gen_nn_ops.conv3d(x, w, strides, 'SAME')
    y = gen_nn_ops.bias_add(y, b, 'NHWC')
    output = array_ops.identity(y)

    with session.Session(config=_get_config(False)) as sess:
      output_val_ref = sess.run(output)

    with session.Session(config=_get_config()) as sess:
      metadata = config_pb2.RunMetadata()
      output_val = sess.run(output, run_metadata=metadata)

    nodes = []
    num_transposes = 0
    for node in metadata.cost_graph.node:
      if _is_transpose(node.name):
        num_transposes += 1
      nodes.append(node.name)

    expected_num_transposes = 2
    self.assertEqual(expected_num_transposes, num_transposes)
    self._assert_trans_ndhwc_to_ncdhw('Conv3D-0', nodes)
    self._assert_trans_ncdhw_to_ndhwc('BiasAdd-0-0', nodes)
    self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testBiasAddGradFor5DTensor(self):
    if not test.is_gpu_available(cuda_only=True):
      self.skipTest('GPU required')
    random_seed.set_random_seed(0)
    dy = random_ops.truncated_normal([2, 2, 14, 14, 1], seed=0)
    w = random_ops.truncated_normal([2, 2, 2, 1, 1], seed=0)
    strides = [1, 1, 1, 1, 1]
    dy_shape = array_ops.shape(dy)
    dx = gen_nn_ops.conv3d_backprop_input_v2(dy_shape, w, dy, strides, 'SAME')
    db = gen_nn_ops.bias_add_grad(dx, 'NHWC')
    output = array_ops.identity(db)

    with session.Session(config=_get_config(False)) as sess:
      output_val_ref = sess.run(output)

    with session.Session(config=_get_config()) as sess:
      metadata = config_pb2.RunMetadata()
      output_val = sess.run(output, run_metadata=metadata)

    nodes = []
    num_transposes = 0
    for node in metadata.cost_graph.node:
      if _is_transpose(node.name):
        num_transposes += 1
      nodes.append(node.name)

    # The output of Conv3DBackpropInputV2 won't be converted back to NDHWC
    # because of the BiasAddGrad.
    expected_num_transposes = 1
    self.assertEqual(expected_num_transposes, num_transposes)
    self._assert_vec_ndhwc_to_ncdhw('Conv3DBackpropInputV2-0', nodes)
    self._assert_trans_ndhwc_to_ncdhw('Conv3DBackpropInputV2-2', nodes)
    self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testSliceWithNonConstAxis(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      size = array_ops.placeholder(dtype='int32')
      s = array_ops.slice(conv, [0, 0, 0, 0], size)
      output = array_ops.identity(s)

      size_val = [1, 2, 3, 4]
      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = sess.run(output, feed_dict={size: size_val})

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(
            output, run_metadata=metadata, feed_dict={
                size: size_val
            })

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Four transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('Slice-0-0', nodes)
      self._assert_vec_nhwc_to_nchw('Slice-2', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testSliceWithNonConstAxisFor5DTensor(self):
    if not test.is_gpu_available(cuda_only=True):
      self.skipTest('GPU required')
    random_seed.set_random_seed(0)
    x = random_ops.truncated_normal([2, 2, 14, 14, 1], seed=0)
    w = random_ops.truncated_normal([2, 2, 2, 1, 2], seed=0)
    strides = [1, 1, 1, 1, 1]
    y = gen_nn_ops.conv3d(x, w, strides, 'SAME')
    size = array_ops.placeholder(dtype='int32')
    s = array_ops.slice(y, [0, 0, 0, 0, 0], size)
    output = array_ops.identity(s)

    size_val = [1, 1, 2, 2, 1]
    with session.Session(config=_get_config(False)) as sess:
      output_val_ref = sess.run(output, feed_dict={size: size_val})

    with session.Session(config=_get_config()) as sess:
      metadata = config_pb2.RunMetadata()
      output_val = sess.run(
          output, run_metadata=metadata, feed_dict={size: size_val})

    nodes = []
    num_transposes = 0
    for node in metadata.cost_graph.node:
      if _is_transpose(node.name):
        num_transposes += 1
      nodes.append(node.name)

    # Four transposes were initially added in the Expand phase of
    # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
    expected_num_transposes = 2
    self.assertEqual(expected_num_transposes, num_transposes)
    self._assert_trans_ndhwc_to_ncdhw('Conv3D-0', nodes)
    self._assert_trans_ncdhw_to_ndhwc('Slice-0-0', nodes)
    self._assert_vec_ndhwc_to_ncdhw('Slice-2', nodes)
    self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testStridedSliceWithNonConstAxis(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      end = array_ops.placeholder(dtype='int32')
      s = array_ops.strided_slice(conv, [0, 0, 0, 0], end, strides=[1, 2, 3, 1])
      output = array_ops.identity(s)

      end_val = [1, 2, 3, 4]
      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = sess.run(output, feed_dict={end: end_val})

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(
            output, run_metadata=metadata, feed_dict={
                end: end_val
            })

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Four transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('StridedSlice-0-0', nodes)
      self._assert_vec_nhwc_to_nchw('StridedSlice-2', nodes)
      self.assertIn('StridedSlice-1-LayoutOptimizer', nodes)
      self.assertIn('StridedSlice-3-LayoutOptimizer', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testStridedSliceWithMask1011(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      # This will generate a StridedSlice op with begin mask and
      # end mask 11(1011).
      s = conv[:, :, 1:-1, :]
      output = array_ops.identity(s)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = self.evaluate(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Four transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('strided_slice-0-0', nodes)
      self.assertIn('strided_slice-1-LayoutOptimizer', nodes)
      self.assertIn('strided_slice-2-LayoutOptimizer', nodes)
      self.assertIn('strided_slice-3-LayoutOptimizer', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testStridedSliceWithMask0111(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      # This will generate a StridedSlice op with begin mask and
      # end mask 7(0111).
      s = conv[:, :, :, 1:-1]
      output = array_ops.identity(s)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = self.evaluate(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Four transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('strided_slice-0-0', nodes)
      self.assertIn('strided_slice-1-LayoutOptimizer', nodes)
      self.assertIn('strided_slice-2-LayoutOptimizer', nodes)
      self.assertIn('strided_slice-3-LayoutOptimizer', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testStridedSliceGradWithNonConstAxis(self):
    if test.is_gpu_available(cuda_only=True):
      random_seed.set_random_seed(0)
      x = random_ops.truncated_normal([1, 784], seed=0)
      conv = _two_layer_model(x)
      end = array_ops.placeholder(dtype='int32')
      shape = array_ops.shape(conv)
      end_val = [1, 2, 3, 4]
      s = array_ops.strided_slice(
          conv, [0, 0, 0, 0], end_val, strides=[1, 2, 3, 1])
      s_grad = array_ops.strided_slice_grad(shape, [0, 0, 0, 0], end,
                                            [1, 2, 3, 1], s)
      output = array_ops.identity(s_grad)

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = sess.run(output, feed_dict={end: end_val})

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(
            output, run_metadata=metadata, feed_dict={
                end: end_val
            })

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Four transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('StridedSliceGrad-0-0', nodes)
      self._assert_vec_nhwc_to_nchw('StridedSliceGrad-2', nodes)
      self.assertIn('StridedSlice-1-LayoutOptimizer', nodes)
      self.assertIn('StridedSlice-2-LayoutOptimizer', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testShapeN(self):
    if test.is_gpu_available(cuda_only=True):
      x = array_ops.placeholder(dtype='float32')
      conv = _two_layer_model(x)
      shapen = array_ops.shape_n([conv, conv])
      output = math_ops.add(shapen[0], shapen[1])

      x_val = [1.7] * 784
      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = sess.run(output, feed_dict={x: x_val})

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(
            output, run_metadata=metadata, feed_dict={
                x: x_val
            })

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      expected_num_transposes = 1
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self._assert_vec_nchw_to_nhwc('ShapeN-0-0', nodes)
      self.assertAllEqual(output_val_ref, output_val)

  @test_util.deprecated_graph_mode_only
  def testShapeNFor5DTensor(self):
    if not test.is_gpu_available(cuda_only=True):
      self.skipTest('GPU required')
    h = array_ops.placeholder(dtype='float32')
    x = array_ops.reshape(h, [-1, 2, 14, 14, 1])
    w = random_ops.truncated_normal([2, 2, 2, 1, 2], seed=0)
    strides = [1, 1, 1, 1, 1]
    y = gen_nn_ops.conv3d(x, w, strides, 'SAME')
    shapen = array_ops.shape_n([y, y])
    output = math_ops.add(shapen[0], shapen[1])

    x_val = [1.7] * 784
    with session.Session(config=_get_config(False)) as sess:
      output_val_ref = sess.run(output, feed_dict={h: x_val})

    with session.Session(config=_get_config()) as sess:
      metadata = config_pb2.RunMetadata()
      output_val = sess.run(output, run_metadata=metadata, feed_dict={h: x_val})

    nodes = []
    num_transposes = 0
    for node in metadata.cost_graph.node:
      if _is_transpose(node.name):
        num_transposes += 1
      nodes.append(node.name)

    expected_num_transposes = 1
    self.assertEqual(expected_num_transposes, num_transposes)
    self._assert_trans_ndhwc_to_ncdhw('Conv3D-0', nodes)
    self._assert_vec_ncdhw_to_ndhwc('ShapeN-0-0', nodes)
    self._assert_vec_ncdhw_to_ndhwc('ShapeN-1-0', nodes)
    self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testIdentityNFor4DAnd5DTensors(self):
    if not test.is_gpu_available(cuda_only=True):
      self.skipTest('GPU required')
    h = array_ops.placeholder(dtype='float32')
    x = array_ops.reshape(h, [-1, 2, 14, 14, 1])
    w = random_ops.truncated_normal([2, 2, 2, 1, 4], seed=0)
    strides = [1, 1, 1, 1, 1]
    y = gen_nn_ops.conv3d(x, w, strides, 'SAME')
    x1 = array_ops.reshape(h, [-1, 784])
    y1 = _two_layer_model(x1)
    outputs = array_ops.identity_n([y1, y])
    new_x0 = array_ops.reshape(outputs[0], [-1, 2, 14, 14, 1])
    new_x1 = array_ops.reshape(outputs[1], [-1, 2, 14, 14, 1])
    output = math_ops.add(new_x0, new_x1)

    x_val = [1.7] * 784
    with session.Session(config=_get_config(False)) as sess:
      output_val_ref = sess.run(output, feed_dict={h: x_val})

    with session.Session(config=_get_config()) as sess:
      metadata = config_pb2.RunMetadata()
      output_val = sess.run(output, run_metadata=metadata, feed_dict={h: x_val})

    nodes = []
    num_transposes = 0
    for node in metadata.cost_graph.node:
      if _is_transpose(node.name):
        num_transposes += 1
      nodes.append(node.name)

    expected_num_transposes = 4
    self.assertEqual(expected_num_transposes, num_transposes)
    self._assert_trans_ndhwc_to_ncdhw('Conv3D-0', nodes)
    self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
    self._assert_trans_ncdhw_to_ndhwc('IdentityN-1-0', nodes)
    self._assert_trans_nchw_to_nhwc('IdentityN-0-0', nodes)
    self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testShapeNFollowedByNotConvertibleNodeReshape(self):
    if test.is_gpu_available(cuda_only=True):
      x = array_ops.placeholder(dtype='float32')
      conv = _two_layer_model(x)
      conv_reshape = array_ops.reshape(conv, [1, 1, 1, -1])
      shapen = array_ops.shape_n([conv, conv_reshape])
      shape = array_ops.identity(shapen[1])
      ones = array_ops.ones(shape)
      output = math_ops.add_n([conv_reshape, ones])

      x_val = [1.7] * 784
      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = sess.run(output, feed_dict={x: x_val})

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(
            output, run_metadata=metadata, feed_dict={x: x_val})

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('Conv2D-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testLoop(self):
    if test.is_gpu_available(cuda_only=True):
      output = _loop()

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = self.evaluate(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      # Four transposes were initially added in the Expand phase of
      # LayoutOptimizer; two of them are cancelled out in the Collapse phase.
      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('map/while/Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('map/while/MaxPool_1-0-2', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testLoopWithBranch(self):
    if test.is_gpu_available(cuda_only=True):
      output = _loop_with_branch()

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = self.evaluate(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      expected_num_transposes = 3
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('map/while/Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('map/while/Add_1-0-2', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testLoopWithVecAnd4D(self):
    if test.is_gpu_available(cuda_only=True):
      output = _loop_with_vec_and_4d()

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = self.evaluate(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('map/while/Conv2D-0', nodes)
      self._assert_trans_nchw_to_nhwc('map/while/Add_1-0-2', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testBinaryOpSecondPort(self):
    if test.is_gpu_available(cuda_only=True):
      output = _model_with_second_port()

      with session.Session(config=_get_config(False)) as sess:
        output_val_ref = self.evaluate(output)

      with session.Session(config=_get_config()) as sess:
        metadata = config_pb2.RunMetadata()
        output_val = sess.run(output, run_metadata=metadata)

      nodes = []
      num_transposes = 0
      for node in metadata.cost_graph.node:
        if _is_transpose(node.name):
          num_transposes += 1
        nodes.append(node.name)

      expected_num_transposes = 2
      self.assertEqual(expected_num_transposes, num_transposes)
      self._assert_trans_nhwc_to_nchw('FusedBatchNormV3-0', nodes)
      self._assert_trans_nchw_to_nhwc('Add-0-0', nodes)
      self.assertAllClose(output_val_ref, output_val, atol=1e-3)

  @test_util.deprecated_graph_mode_only
  def testGradient(self):
    meta_graph = _simple_metagraph()
    config = config_pb2.ConfigProto()
    config.graph_options.rewrite_options.CopyFrom(
        rewriter_config_pb2.RewriterConfig(
            layout_optimizer=rewriter_config_pb2.RewriterConfig.ON,
            min_graph_nodes=-1))
    optimized_graph = tf_optimizer.OptimizeGraph(
        config, meta_graph, cluster=_get_cluster())

    found = 0
    for node in optimized_graph.node:
      if node.op in ['Conv2D', 'Conv2DBackpropFilter', 'Conv2DBackpropInput']:
        found += 1
        self.assertEqual(node.attr['data_format'].s, b'NCHW')
    self.assertEqual(found, 5)

  @test_util.deprecated_graph_mode_only
  def testDepthwise(self):
    meta_graph = _simple_metagraph(depthwise=True)
    config = config_pb2.ConfigProto()
    config.graph_options.rewrite_options.CopyFrom(
        rewriter_config_pb2.RewriterConfig(
            layout_optimizer=rewriter_config_pb2.RewriterConfig.ON,
            min_graph_nodes=-1))
    optimized_graph = tf_optimizer.OptimizeGraph(
        config, meta_graph, cluster=_get_cluster())

    found = 0
    for node in optimized_graph.node:
      if node.op in [
          'DepthwiseConv2dNative', 'DepthwiseConv2dNativeBackpropFilter',
          'DepthwiseConv2dNativeBackpropInput'
      ]:
        found += 1
        self.assertEqual(node.attr['data_format'].s, b'NCHW')
    self.assertEqual(found, 6)

  def testCheckpointCompatibility(self):
    if not test.is_gpu_available(cuda_only=True):
      self.skipTest('GPU required')

    checkpoint_path = self.get_temp_dir()
    self._train(checkpoint_path)
    vars_expected = self._train(checkpoint_path, restore=True)
    vars_layout_optimized = self._train(
        checkpoint_path, restore=True, layout_optimizer=True)

    for var_expected, var_layout_optimized in zip(vars_expected,
                                                  vars_layout_optimized):
      self.assertAllClose(var_expected, var_layout_optimized, atol=1e-6)


if __name__ == '__main__':
  test.main()