# Copyright 2016 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 Cudnn RNN models."""
from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import argparse
import collections
import functools
import itertools
import os
import sys
import unittest

import numpy as np

from tensorflow.contrib.cudnn_rnn.python.layers import cudnn_rnn
from tensorflow.contrib.cudnn_rnn.python.ops import cudnn_rnn_ops
from tensorflow.contrib.rnn.python.ops import rnn as contrib_rnn_lib
from tensorflow.python.eager import backprop
from tensorflow.python.eager import context
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.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import gen_nn_ops
from tensorflow.python.ops import gradients_impl as gradients
from tensorflow.python.ops import init_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.ops import rnn as rnn_lib
from tensorflow.python.ops import rnn_cell_impl
from tensorflow.python.ops import state_ops
from tensorflow.python.ops import variable_scope as vs
from tensorflow.python.ops import variables
from tensorflow.python.ops.losses import losses
from tensorflow.python.platform import googletest
from tensorflow.python.platform import test
from tensorflow.python.platform import tf_logging as logging
from tensorflow.python.training import adagrad
from tensorflow.python.training import adam
from tensorflow.python.training import gradient_descent
from tensorflow.python.training import momentum
from tensorflow.python.training import rmsprop
from tensorflow.python.training import saver as saver_lib
from tensorflow.python.training.checkpointable import util as checkpointable_utils


CUDNN_LSTM = cudnn_rnn_ops.CUDNN_LSTM
CUDNN_GRU = cudnn_rnn_ops.CUDNN_GRU
CUDNN_RNN_RELU = cudnn_rnn_ops.CUDNN_RNN_RELU
CUDNN_RNN_TANH = cudnn_rnn_ops.CUDNN_RNN_TANH
CUDNN_RNN_UNIDIRECTION = cudnn_rnn_ops.CUDNN_RNN_UNIDIRECTION
CUDNN_RNN_BIDIRECTION = cudnn_rnn_ops.CUDNN_RNN_BIDIRECTION

CUDNN_LSTM_PARAMS_PER_LAYER = cudnn_rnn_ops.CUDNN_LSTM_PARAMS_PER_LAYER
CUDNN_GRU_PARAMS_PER_LAYER = cudnn_rnn_ops.CUDNN_GRU_PARAMS_PER_LAYER
CUDNN_RNN_TANH_PARAMS_PER_LAYER = cudnn_rnn_ops.CUDNN_RNN_TANH_PARAMS_PER_LAYER
CUDNN_RNN_RELU_PARAMS_PER_LAYER = cudnn_rnn_ops.CUDNN_RNN_RELU_PARAMS_PER_LAYER


class CudnnTestModel(object):
  """Model with convenient APIs for easier building and running test graph.

  The graph built is used by all tests below to avoid repeatedly building
  similar test graphs.
  """

  def __init__(self,
               rnn_mode,
               num_layers,
               num_units,
               input_size,
               direction=CUDNN_RNN_UNIDIRECTION,
               dropout=0.,
               dtype=dtypes.float32,
               training=False,
               seed=None,
               kernel_initializer=None,
               bias_initializer=None):
    if dtype not in (dtypes.float16, dtypes.float32, dtypes.float64):
      raise ValueError("Invalid dtype: %s" % dtype)
    self._dtype = dtype

    self._inputs = array_ops.placeholder(
        dtype=dtype, shape=[None, None, input_size], name="inputs")
    h = array_ops.placeholder(
        dtype=dtype, shape=[None, None, num_units], name="h")
    c = array_ops.placeholder(
        dtype=dtype, shape=[None, None, num_units], name="c")
    if rnn_mode == CUDNN_LSTM:
      model_fn = cudnn_rnn.CudnnLSTM
      self._initial_state = (h, c)
    elif rnn_mode == CUDNN_GRU:
      model_fn = cudnn_rnn.CudnnGRU
      self._initial_state = (h,)
    elif rnn_mode == CUDNN_RNN_TANH:
      model_fn = cudnn_rnn.CudnnRNNTanh
      self._initial_state = (h,)
    elif rnn_mode == CUDNN_RNN_RELU:
      model_fn = cudnn_rnn.CudnnRNNRelu
      self._initial_state = (h,)
    else:
      raise ValueError("Invalid rnn_mode: %s" % rnn_mode)
    self._rnn = model_fn(
        num_layers,
        num_units,
        direction=direction,
        dropout=dropout,
        dtype=dtype,
        seed=seed,
        kernel_initializer=kernel_initializer,
        bias_initializer=bias_initializer)
    self._rnn.build([None, None, input_size])

    self._outputs, self._output_state = self._rnn(
        self._inputs, initial_state=self._initial_state, training=training)

  def _AddUp(self, outputs, output_state):
    total = math_ops.reduce_sum(outputs)
    for s in output_state:
      total += math_ops.reduce_sum(s)
    return total

  @property
  def inputs(self):
    return self._inputs

  @property
  def initial_state(self):
    return self._initial_state

  @property
  def outputs(self):
    return self._outputs

  @property
  def output_state(self):
    return self._output_state

  @property
  def rnn(self):
    return self._rnn

  @property
  def total_sum(self):
    return self._AddUp(self.outputs, self.output_state)

  def SynthesizeInput(self, seq_length, batch_size, seed=1234):
    """Synthesizes input and initial state values for testing."""
    np.random.seed(seed)
    num_layers = self._rnn.num_layers
    dir_count = self._rnn.num_dirs
    num_units = self._rnn.num_units
    input_size = self._rnn.input_size

    np_dtype = np.float32 if self._dtype == dtypes.float32 else np.float64
    inputs = np.random.randn(seq_length, batch_size,
                             input_size).astype(np_dtype)
    input_h = np.random.randn(num_layers * dir_count, batch_size,
                              num_units).astype(np_dtype)
    if self._rnn.rnn_mode == CUDNN_LSTM:
      input_c = np.random.randn(num_layers * dir_count, batch_size,
                                num_units).astype(np_dtype)
      initial_state = (input_h, input_c)
    else:
      initial_state = (input_h,)
    return inputs, initial_state

  def ZeroState(self, batch_size):
    num_layers = self._rnn.num_layers
    dir_count = self._rnn.num_dirs
    num_units = self._rnn.num_units

    np_dtype = np.float32 if self._dtype == dtypes.float32 else np.float64
    input_h = np.zeros((num_layers * dir_count, batch_size,
                        num_units)).astype(np_dtype)
    if self._rnn.rnn_mode == CUDNN_LSTM:
      input_c = np.zeros((num_layers * dir_count, batch_size,
                          num_units)).astype(np_dtype)
      initial_state = (input_h, input_c)
    else:
      initial_state = (input_h,)
    return initial_state

  def FProp(self, inputs_t, initial_state_t, training):
    """Builds additional subgraph with given inputs and state.

    Args:
      inputs_t: a tensor.
      initial_state_t: a tensor.
      training: boolean, true if training mode.
    Returns:
      A tensor of the forward pass output of the model.
    """
    outputs, output_state = self._rnn(
        inputs_t, initial_state=initial_state_t, training=training)
    return self._AddUp(outputs, output_state)

  def Feed(self, sess, inputs, initial_state=None, return_sum=True):
    """Runs graph with given inputs and initial state."""
    batch_size = inputs.shape[1]
    if initial_state is None:
      initial_state = self.ZeroState(batch_size)
    if return_sum:
      return sess.run(
          self.total_sum,
          feed_dict={self.inputs: inputs,
                     self.initial_state: initial_state})
    else:
      return sess.run(
          [self.outputs, self.output_state],
          feed_dict={self.inputs: inputs,
                     self.initial_state: initial_state})


def _CreateCudnnCompatibleCanonicalRNN(rnn, inputs, is_bidi=False, scope=None):
  mode = rnn.rnn_mode
  num_units = rnn.num_units
  num_layers = rnn.num_layers

  # To reuse cuDNN-trained models, must use cudnn compatible rnn cells.
  if mode == CUDNN_LSTM:
    single_cell = lambda: cudnn_rnn_ops.CudnnCompatibleLSTMCell(num_units)
  elif mode == CUDNN_GRU:
    single_cell = lambda: cudnn_rnn_ops.CudnnCompatibleGRUCell(num_units)
  elif mode == CUDNN_RNN_TANH:
    single_cell = (lambda: rnn_cell_impl.BasicRNNCell(num_units, math_ops.tanh))
  elif mode == CUDNN_RNN_RELU:
    single_cell = (
        lambda: rnn_cell_impl.BasicRNNCell(num_units, gen_nn_ops.relu))
  else:
    raise ValueError("%s is not supported!" % mode)

  if not is_bidi:
    cell = rnn_cell_impl.MultiRNNCell(
        [single_cell() for _ in range(num_layers)])
    return rnn_lib.dynamic_rnn(
        cell, inputs, dtype=dtypes.float32, time_major=True, scope=scope)
  else:
    cells_fw = [single_cell() for _ in range(num_layers)]
    cells_bw = [single_cell() for _ in range(num_layers)]

    (outputs, output_state_fw,
     output_state_bw) = contrib_rnn_lib.stack_bidirectional_dynamic_rnn(
         cells_fw,
         cells_bw,
         inputs,
         dtype=dtypes.float32,
         time_major=True,
         scope=scope)
    return outputs, (output_state_fw, output_state_bw)


class CudnnRNNTestBasic(test_util.TensorFlowTestCase):

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testLayerBasic(self):
    num_layers = 4
    num_units = 2
    batch_size = 8
    direction = CUDNN_RNN_UNIDIRECTION
    dir_count = 1

    with vs.variable_scope("main"):
      kernel_initializer = init_ops.constant_initializer(0.)
      bias_initializer = init_ops.constant_initializer(0.)
      inputs = random_ops.random_uniform([
          num_layers * dir_count, batch_size, num_units], dtype=dtypes.float32)

      lstm = cudnn_rnn.CudnnLSTM(num_layers, num_units,
                                 direction=direction,
                                 kernel_initializer=kernel_initializer,
                                 bias_initializer=bias_initializer,
                                 name="awesome_lstm")

      # Build the layer
      outputs1, _ = lstm(inputs)
      # Reuse the layer
      outputs2, _ = lstm(inputs)

      total_sum1 = math_ops.reduce_sum(outputs1)
      total_sum2 = math_ops.reduce_sum(outputs2)

    with vs.variable_scope("main", reuse=True):
      lstm = cudnn_rnn.CudnnLSTM(num_layers, num_units,
                                 direction=direction,
                                 kernel_initializer=kernel_initializer,
                                 bias_initializer=bias_initializer,
                                 name="awesome_lstm")

      # Reuse the layer
      outputs3, _ = lstm(inputs)
      total_sum3 = math_ops.reduce_sum(outputs3)

    self.assertEqual(1, len(variables.trainable_variables()))
    self.assertEqual(1, len(ops.get_collection(ops.GraphKeys.SAVEABLE_OBJECTS)))
    self.assertEqual("main/awesome_lstm/opaque_kernel",
                     variables.trainable_variables()[0].op.name)

    with self.test_session(use_gpu=True) as sess:
      sess.run(variables.global_variables_initializer())
      (total_sum1_v, total_sum2_v, total_sum3_v) = sess.run(
          [total_sum1, total_sum2, total_sum3])
      self.assertEqual(0, total_sum1_v)
      self.assertEqual(0, total_sum2_v)
      self.assertEqual(0, total_sum3_v)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testOptimizersSupport(self):
    for opt in ("adagrad", "adam", "rmsprop", "momentum", "sgd"):
      self._TestOptimizerSupportHelper(opt)

  def _GetOptimizer(self, opt):
    if opt == "adagrad":
      return adagrad.AdagradOptimizer(learning_rate=1e-2)
    elif opt == "adam":
      return adam.AdamOptimizer(learning_rate=1e-2)
    elif opt == "rmsprop":
      return rmsprop.RMSPropOptimizer(learning_rate=1e-2)
    elif opt == "momentum":
      return momentum.MomentumOptimizer(learning_rate=1e-2, momentum=0.9)
    elif opt == "sgd":
      return gradient_descent.GradientDescentOptimizer(learning_rate=1e-2)
    else:
      raise ValueError("Unsupported optimizer: %s" % opt)

  def _TestOptimizerSupportHelper(self, opt):
    num_layers = 4
    num_units = 2
    batch_size = 8
    direction = CUDNN_RNN_UNIDIRECTION
    dir_count = 1

    with ops.Graph().as_default() as g:
      kernel_initializer = init_ops.constant_initializer(0.)
      bias_initializer = init_ops.constant_initializer(0.)
      inputs = random_ops.random_uniform([
          num_layers * dir_count, batch_size, num_units], dtype=dtypes.float32)

      lstm = cudnn_rnn.CudnnLSTM(num_layers, num_units,
                                 direction=direction,
                                 kernel_initializer=kernel_initializer,
                                 bias_initializer=bias_initializer,
                                 name="awesome_lstm")
      outputs, _ = lstm(inputs)
      loss = math_ops.reduce_sum(outputs)
      optimizer = self._GetOptimizer(opt)
      train_op = optimizer.minimize(loss)

    with self.test_session(use_gpu=True, graph=g) as sess:
      sess.run(variables.global_variables_initializer())
      sess.run(train_op)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testSaveableGraphDeviceAssignment(self):
    num_layers = 4
    num_units = 2
    batch_size = 8
    direction = CUDNN_RNN_UNIDIRECTION
    dir_count = 1

    def DeviceFn(op):
      if op.type in ("Variable", "VariableV2"):
        return "/cpu:0"
      else:
        return "/gpu:0"

    with ops.Graph().as_default() as g:
      with ops.device(DeviceFn):
        with vs.variable_scope("main"):
          kernel_initializer = init_ops.constant_initializer(3.14)
          bias_initializer = init_ops.constant_initializer(1.59)
          inputs = random_ops.random_uniform(
              [num_layers * dir_count, batch_size, num_units],
              dtype=dtypes.float32)

          lstm = cudnn_rnn.CudnnLSTM(num_layers, num_units,
                                     direction=direction,
                                     kernel_initializer=kernel_initializer,
                                     bias_initializer=bias_initializer,
                                     name="awesome_lstm")
          outputs = lstm(inputs)

        # saver is created in the scope of DeviceFn.
        saver = saver_lib.Saver()

    with self.test_session(use_gpu=True, graph=g) as sess:
      save_path = os.path.join(self.get_temp_dir(),
                               "test-saveable-device-assignment")
      sess.run(variables.global_variables_initializer())

      saver.save(sess, save_path)
      saver.restore(sess, save_path)
      sess.run(outputs)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testDifferentShapesEager(self):
    # Checks that kernel caching does not cause sharing of temporary storage
    # across different input shapes when executing eagerly.
    with context.eager_mode():
      with ops.device("gpu:0"):
        first_output, _ = cudnn_rnn.CudnnGRU(1, 100)(
            array_ops.zeros([28, 100, 28]))
        second_output, _ = cudnn_rnn.CudnnGRU(1, 100)(
            array_ops.zeros([28, 100, 100]))
        self.assertAllEqual([28, 100, 100], first_output.shape)
        self.assertAllEqual([28, 100, 100], second_output.shape)

        def _LossFunc():
          first_output, _ = cudnn_rnn.CudnnGRU(1, 100)(
              array_ops.zeros([28, 100, 28]))
          second_output, _ = cudnn_rnn.CudnnGRU(1, 100)(
              array_ops.zeros([28, 100, 100]))
          return (math_ops.reduce_sum(first_output) +
                  math_ops.reduce_sum(second_output))

        backprop.implicit_grad(_LossFunc)()

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testDifferentShapesGraph(self):
    # Tests that a single kernel instance presented with multiple input shapes
    # does not crash with graph execution.
    with ops.device("gpu:0"):
      layer = cudnn_rnn.CudnnGRU(1, 100)
      layer(array_ops.zeros([28, 100, 100]))

      def _Cond(index, accumulation):
        del accumulation  # unused
        return math_ops.less(index, 4)

      def _Body(index, accumulation):
        layer_input = accumulation[:, :, 10 * (1 + index % 2):]
        output, _ = layer(layer_input)
        return index + 1, accumulation + output

      original_input = array_ops.zeros([28, 100, 100])
      _, accumulation = control_flow_ops.while_loop(_Cond, _Body,
                                                    [0, original_input])
      grad, = gradients.gradients(
          math_ops.reduce_sum(accumulation), (original_input,))
    init_op = variables.global_variables_initializer()
    with self.cached_session() as sess:
      sess.run(init_op)
      accumulation_eval, grad_eval = sess.run((accumulation, grad))
      self.assertAllEqual([28, 100, 100], accumulation_eval.shape)
      self.assertAllEqual([28, 100, 100], grad_eval.shape)


# TODO(jamesqin): Transform to parameterized test after it is included in the
# TF open source codebase.
class CudnnRNNTestSaveRestore(test_util.TensorFlowTestCase):

  def _CompareWeights(self, lhs, rhs):
    self.assertEqual(len(lhs), len(rhs))
    for lw, rw in zip(lhs, rhs):
      self.assertAllEqual(lw, rw)

  def _CompareBiases(self, lhs, rhs, rnn_mode, num_layers, direction):
    self.assertEqual(len(lhs), len(rhs))
    if rnn_mode == CUDNN_LSTM:
      num_params_per_layer = CUDNN_LSTM_PARAMS_PER_LAYER
    elif rnn_mode == CUDNN_GRU:
      num_params_per_layer = CUDNN_GRU_PARAMS_PER_LAYER
    elif rnn_mode == CUDNN_RNN_TANH:
      num_params_per_layer = CUDNN_RNN_TANH_PARAMS_PER_LAYER
    else:
      num_params_per_layer = CUDNN_RNN_RELU_PARAMS_PER_LAYER
    num_dirs = 1 if direction == CUDNN_RNN_UNIDIRECTION else 2
    num_params_per_layer *= num_dirs
    self.assertEqual(num_params_per_layer * num_layers, len(lhs))

    for i in range(num_layers):
      layer_lhs = lhs[i * num_params_per_layer: (i+1) * num_params_per_layer]
      layer_rhs = rhs[i * num_params_per_layer: (i+1) * num_params_per_layer]
      if direction == CUDNN_RNN_UNIDIRECTION:
        self._CompareSingleLayerBiases(layer_lhs, layer_rhs)
      else:
        size = len(layer_lhs)
        fw_lhs, bw_lhs = layer_lhs[:size//2], layer_lhs[size//2:]
        fw_rhs, bw_rhs = layer_rhs[:size//2], layer_rhs[size//2:]
        self._CompareSingleLayerBiases(fw_lhs, fw_rhs)
        self._CompareSingleLayerBiases(bw_lhs, bw_rhs)

  def _CompareSingleLayerBiases(self, lhs, rhs):
    self.assertEqual(len(lhs), len(rhs))

    lf_lhs, rt_lhs = lhs[:len(lhs)//2], lhs[len(lhs)//2:]
    lf_rhs, rt_rhs = rhs[:len(rhs)//2], rhs[len(rhs)//2:]
    self.assertEqual(len(lf_lhs), len(rt_lhs))
    self.assertEqual(len(lf_rhs), len(rt_rhs))

    sum_lhs, sum_rhs = [], []
    for lf, rt in zip(lf_lhs, rt_lhs):
      sum_lhs.append(lf + rt)
    for lf, rt in zip(lf_rhs, rt_rhs):
      sum_rhs.append(lf + rt)
    self.assertEqual(len(sum_lhs), len(sum_rhs))
    for lf, rt in zip(sum_lhs, sum_rhs):
      self.assertAllEqual(lf, rt)

  def _TestSaveRestoreVariable(self, rnn_mode, direction, dtype):
    input_size = 3
    num_layers = 2
    num_units = 7
    with ops.Graph().as_default() as g:
      random_seed.set_random_seed(1234)
      model = CudnnTestModel(
          rnn_mode,
          num_layers,
          num_units,
          input_size,
          direction=direction,
          dtype=dtype)
      rnn = model.rnn
      save_path = os.path.join(self.get_temp_dir(),
                               "save-restore-variable-test")
      saver = saver_lib.Saver()
      weights, biases = (
          model.rnn.saveable.format_converter._opaque_to_cu_canonical(
              model.rnn.saveable._variables))
      opaque_params = rnn.trainable_variables[0]
      # CudnnTestModel() creates CudnnOpaqueParamsSaveable that helps saver save
      # Cudnn vars in canonical format.
      reset_op = state_ops.assign(
          opaque_params,
          array_ops.zeros(array_ops.shape(opaque_params), dtype=dtype))
      # Passing graph explicitly, otherwise an old sess would be reused.
      with self.test_session(use_gpu=True, graph=g) as sess:
        sess.run(variables.global_variables_initializer())
        val = saver.save(sess, save_path)
        self.assertEqual(save_path, val)
        weights_v, biases_v = sess.run([weights, biases])

        # Reset opaque param
        sess.run(reset_op)
        saver.restore(sess, save_path)
        weights_v_restored, biases_v_restored = sess.run([weights, biases])

        self._CompareWeights(weights_v, weights_v_restored)
        self._CompareBiases(biases_v, biases_v_restored, rnn_mode, num_layers,
                            direction)

  def _TestSaveRestoreTwoVariables(self, rnn_mode, direction, dtype):
    input_size = 3
    num_layers = 2
    num_units = 7
    with ops.Graph().as_default() as g:
      random_seed.set_random_seed(1234)
      with vs.variable_scope("m1"):
        model1 = CudnnTestModel(
            rnn_mode,
            num_layers,
            num_units,
            input_size,
            direction=direction,
            dtype=dtype)
      with vs.variable_scope("m2"):
        model2 = CudnnTestModel(
            rnn_mode,
            num_layers,
            num_units,
            input_size,
            direction=direction,
            dtype=dtype)
      opaque_params = (model1.rnn.trainable_variables[0],
                       model2.rnn.trainable_variables[0])
      saveable1 = model1.rnn.saveable
      weights1, biases1 = saveable1.format_converter._opaque_to_cu_canonical(
          saveable1._variables)
      saveable2 = model1.rnn.saveable
      weights2, biases2 = saveable2.format_converter._opaque_to_cu_canonical(
          saveable2._variables)
      reset_params = [
          state_ops.assign(params,
                           array_ops.zeros_like(params, dtype=dtype))
          for params in opaque_params
      ]
      reset_op = control_flow_ops.group(*reset_params)
      save_path = os.path.join(self.get_temp_dir(),
                               "save-restore-variable-test2")
      saver = saver_lib.Saver()
      # Passing graph explicitly, otherwise an old sess would be reused.
      with self.test_session(use_gpu=True, graph=g) as sess:
        sess.run(variables.global_variables_initializer())
        val = saver.save(sess, save_path)
        self.assertEqual(save_path, val)

        weights1_v, biases1_v = sess.run([weights1, biases1])
        weights2_v, biases2_v = sess.run([weights2, biases2])

        sess.run(reset_op)
        saver.restore(sess, save_path)
        weights1_v_restored, biases1_v_restored = sess.run([weights1, biases1])
        weights2_v_restored, biases2_v_restored = sess.run([weights2, biases2])

        self._CompareWeights(weights1_v, weights1_v_restored)
        self._CompareWeights(weights2_v, weights2_v_restored)
        self._CompareBiases(biases1_v, biases1_v_restored, rnn_mode, num_layers,
                            direction)
        self._CompareBiases(biases2_v, biases2_v_restored, rnn_mode, num_layers,
                            direction)

  def _TestSaveRestoreOutput(self, rnn_mode, direction, dtype):
    with ops.Graph().as_default() as g:
      num_layers = 2
      num_units = 7
      input_size = 7
      seq_length = 8
      batch_size = 4
      model = CudnnTestModel(
          rnn_mode,
          num_layers,
          num_units,
          input_size,
          direction=direction,
          dtype=dtype,
          training=False)
      rnn = model.rnn

      save_path = os.path.join(self.get_temp_dir(), "save-restore-output-test")
      saver = saver_lib.Saver()

      # Only one opaque var in a cudnn layer.
      assert len(rnn.trainable_variables) == 1
      reset_params = state_ops.assign(
          rnn.trainable_variables[0],
          array_ops.zeros(
              array_ops.shape(rnn.trainable_variables[0]), dtype=dtype))

      # Passing graph explicitly, otherwise an old sess would be reused.
      with self.test_session(use_gpu=True, graph=g) as sess:
        sess.run(variables.global_variables_initializer())
        inputs, initial_state = model.SynthesizeInput(seq_length, batch_size)
        total_sum_v = model.Feed(sess, inputs, initial_state)
        val = saver.save(sess, save_path)
        self.assertEqual(save_path, val)

        sess.run(reset_params)
        saver.restore(sess, save_path)
        total_sum_v_restored = model.Feed(sess, inputs, initial_state)
        self.assertAllClose(total_sum_v, total_sum_v_restored, atol=1e-5)

  def _TestSaveRestoreHelper(self, rnn_mode):
    directions = [CUDNN_RNN_UNIDIRECTION, CUDNN_RNN_BIDIRECTION]
    dtype_list = [dtypes.float16, dtypes.float32, dtypes.float64]
    for direction, dtype in itertools.product(directions, dtype_list):
      self._TestSaveRestoreVariable(rnn_mode, direction, dtype)
      self._TestSaveRestoreTwoVariables(rnn_mode, direction, dtype)
      self._TestSaveRestoreOutput(rnn_mode, direction, dtype)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testSaveRestoreRepeatedlyCreateCustomSaveable(self):
    input_size = 3
    num_layers = 2
    num_units = 7
    with ops.Graph().as_default():
      random_seed.set_random_seed(1234)
      model = CudnnTestModel(
          CUDNN_LSTM,
          num_layers,
          num_units,
          input_size,
          direction=CUDNN_RNN_UNIDIRECTION,
          dtype=dtypes.float32)
      with self.assertRaisesRegexp(RuntimeError,
                                   "Cudnn saveable already created"):
        model.rnn._create_saveable()

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testSaveRestoreLSTM(self):
    self._TestSaveRestoreHelper(CUDNN_LSTM)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testSaveRestoreGRU(self):
    self._TestSaveRestoreHelper(CUDNN_GRU)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testSaveRestoreRNNTanh(self):
    self._TestSaveRestoreHelper(CUDNN_RNN_TANH)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testSaveRestoreRNNRelu(self):
    self._TestSaveRestoreHelper(CUDNN_RNN_RELU)


class CudnnRNNTestSaveRestoreCheckpointable(test_util.TensorFlowTestCase):

  def _VerifyCheckpoint(
      self, checkpoint_path, compatible_cell_fn, cudnn_cell_fn,
      num_layers, input_size, expected_variable_values, num_applications=3):
    checkpoint_directory = self.get_temp_dir()
    checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
    with ops.device("gpu:0"):
      cudnn_layer = cudnn_cell_fn()
      cudnn_checkpoint = checkpointable_utils.Checkpoint(cell=cudnn_layer)
      status = cudnn_checkpoint.restore(checkpoint_path)
      inputs = 3. * array_ops.ones([num_applications, num_layers, input_size],
                                   dtype=dtypes.float32)
      cudnn_output, _ = cudnn_layer(inputs)
      status.run_restore_ops()
    second_save_path = cudnn_checkpoint.save(checkpoint_prefix)
    restore_layer = compatible_cell_fn()
    restore_layer_checkpoint = checkpointable_utils.Checkpoint(
        cell=restore_layer)
    status = restore_layer_checkpoint.restore(second_save_path)
    current_state = restore_layer.zero_state(1, dtypes.float32)
    for _ in range(num_applications):
      restore_layer_output, current_state = restore_layer(
          inputs=3. * array_ops.ones([1, input_size]),
          state=current_state)
    status.run_restore_ops()
    self.assertTrue(restore_layer.variables)
    for variable, expected_value in zip(
        restore_layer.variables, expected_variable_values):
      self.assertAllClose(expected_value, self.evaluate(variable))
    self.assertAllClose(self.evaluate(restore_layer_output),
                        self.evaluate(cudnn_output)[-1, -1:, ...])

  def _CheckpointableSingleCellUnidirectionalTestTemplate(
      self, single_cell_fn, cudnn_cell_fn):
    # Single-layer cuDNN cells with object-based checkpointing should be
    # checkpoint compatible with either single CudnnCompatible cells or
    # MultiRnnCells with one cell.
    input_size = 3
    save_cell_layer = single_cell_fn()
    save_cell_layer(
        inputs=array_ops.ones([1, input_size]),
        state=save_cell_layer.zero_state(1, dtypes.float32))
    self.assertTrue(save_cell_layer.variables)
    expected_values = []
    np.random.seed(10)
    for variable in save_cell_layer.variables:
      value = np.random.normal(size=variable.shape)
      expected_values.append(value)
      self.evaluate(variable.assign(value))
    save_checkpoint = checkpointable_utils.Checkpoint(cell=save_cell_layer)
    checkpoint_directory = self.get_temp_dir()
    checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
    first_save_path = save_checkpoint.save(checkpoint_prefix)
    self._VerifyCheckpoint(
        checkpoint_path=first_save_path,
        compatible_cell_fn=
        lambda: rnn_cell_impl.MultiRNNCell([single_cell_fn()]),
        cudnn_cell_fn=cudnn_cell_fn,
        num_layers=1,
        expected_variable_values=expected_values,
        input_size=input_size)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  @test_util.run_in_graph_and_eager_modes
  def testLSTMCheckpointableSingleLayer(self):
    num_units = 2
    direction = CUDNN_RNN_UNIDIRECTION
    self._CheckpointableSingleCellUnidirectionalTestTemplate(
        single_cell_fn=functools.partial(
            cudnn_rnn_ops.CudnnCompatibleLSTMCell, num_units=num_units),
        cudnn_cell_fn=functools.partial(
            cudnn_rnn.CudnnLSTM, num_layers=1, num_units=num_units,
            direction=direction, name="awesome_lstm"))

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  @test_util.run_in_graph_and_eager_modes
  def testGRUCheckpointableSingleLayer(self):
    num_units = 2
    direction = CUDNN_RNN_UNIDIRECTION
    with self.assertRaises(NotImplementedError):
      # TODO(allenl): Implement object-based saving for GRUs and other cells.
      self._CheckpointableSingleCellUnidirectionalTestTemplate(
          single_cell_fn=functools.partial(
              cudnn_rnn_ops.CudnnCompatibleGRUCell, num_units=num_units),
          cudnn_cell_fn=functools.partial(
              cudnn_rnn.CudnnGRU, num_layers=1, num_units=num_units,
              direction=direction, name="awesome_gru"))

  def _CheckpointableMultiLayerTestTemplate(
      self, single_cell_fn, cudnn_cell_fn, num_layers):

    def _MultiCellFn():
      return rnn_cell_impl.MultiRNNCell(
          [single_cell_fn() for _ in range(num_layers)])
    input_size = 3
    save_graph = ops.Graph()
    with save_graph.as_default(), self.session(graph=save_graph):
      save_layer = _MultiCellFn()
      save_layer(inputs=array_ops.ones([1, input_size]),
                 state=save_layer.zero_state(1, dtypes.float32))
      self.assertTrue(save_layer.variables)
      expected_values = []
      np.random.seed(10)
      for variable in save_layer.variables:
        value = np.random.normal(size=variable.shape)
        expected_values.append(value)
        self.evaluate(variable.assign(value))
      save_checkpoint = checkpointable_utils.Checkpoint(cell=save_layer)
      checkpoint_directory = self.get_temp_dir()
      checkpoint_prefix = os.path.join(checkpoint_directory, "ckpt")
      first_save_path = save_checkpoint.save(checkpoint_prefix)
    self._VerifyCheckpoint(
        checkpoint_path=first_save_path,
        compatible_cell_fn=_MultiCellFn, cudnn_cell_fn=cudnn_cell_fn,
        num_layers=num_layers,
        expected_variable_values=expected_values,
        input_size=input_size)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  @test_util.run_in_graph_and_eager_modes
  def testCudnnCompatibleLSTMCheckpointablMultiLayer(self):
    num_units = 2
    num_layers = 3
    direction = CUDNN_RNN_UNIDIRECTION
    self._CheckpointableMultiLayerTestTemplate(
        single_cell_fn=functools.partial(
            cudnn_rnn_ops.CudnnCompatibleLSTMCell, num_units=num_units),
        cudnn_cell_fn=functools.partial(
            cudnn_rnn.CudnnLSTM, num_layers=num_layers, num_units=num_units,
            direction=direction, name="awesome_lstm"),
        num_layers=num_layers)


# TODO(jamesqin): Transform to parameterized test after it is included in the
# TF open source codebase.
class CudnnRNNTestCompatibleRNNCells(test_util.TensorFlowTestCase):

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testCudnnCompatibleLSTM(self):
    self._TestCudnnCompatibleRnnCellsHelper(CUDNN_LSTM)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testCudnnCompatibleGRU(self):
    self._TestCudnnCompatibleRnnCellsHelper(CUDNN_GRU)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testCudnnCompatibleRNNTanh(self):
    self._TestCudnnCompatibleRnnCellsHelper(CUDNN_RNN_TANH)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testCudnnCompatibleRNNRelu(self):
    self._TestCudnnCompatibleRnnCellsHelper(CUDNN_RNN_RELU)

  def _TestCudnnCompatibleRnnCellsHelper(self, rnn_mode):
    configs = [
        {
            "num_layers": 1,
            "seq_length": 3,
            "num_units": 4,
            "input_size": 5,
            "batch_size": 6,
        },
        {
            "num_layers": 2,
            "seq_length": 8,
            "num_units": 4,
            "input_size": 8,
            "batch_size": 16,
        },
        {
            "num_layers": 2,
            "seq_length": 3,
            "num_units": 4,
            "input_size": 5,
            "batch_size": 6,
        },
        {
            "num_layers": 1,
            "seq_length": 2,
            "num_units": 2,
            "input_size": 4,
            "batch_size": 1,
        },
    ]
    directions = [CUDNN_RNN_UNIDIRECTION, CUDNN_RNN_BIDIRECTION]
    for cfg, direction in zip(configs, directions):
      self._TestCudnnCompatibleRnnCells(cfg["num_layers"], cfg["seq_length"],
                                        cfg["num_units"], cfg["input_size"],
                                        cfg["batch_size"], rnn_mode, direction)

  def _TestCudnnCompatibleRnnCells(self, num_layers, seq_length, num_units,
                                   input_size, batch_size, rnn_mode, direction):
    dtype = dtypes.float32
    # Train graph
    with ops.Graph().as_default() as g:
      model = CudnnTestModel(
          rnn_mode,
          num_layers,
          num_units,
          input_size,
          direction=direction,
          dtype=dtype,
          training=True)
      target_output = array_ops.placeholder(dtype=dtype)
      loss_op = losses.log_loss(
          labels=target_output, predictions=model.total_sum)
      optimizer = gradient_descent.GradientDescentOptimizer(learning_rate=1e-2)
      train_op = optimizer.minimize(loss_op)

      saver = saver_lib.Saver()

      # Train Cudnn model
      seed = 0
      with self.test_session(use_gpu=True, graph=g) as sess:
        sess.run(variables.global_variables_initializer())
        # Train 128 steps
        num_steps = 128
        for _ in range(num_steps):
          inputs, _ = model.SynthesizeInput(seq_length, batch_size, seed)
          targets = np.random.rand()
          sess.run(
              train_op,
              feed_dict={
                  model.inputs: inputs,
                  model.initial_state: model.ZeroState(batch_size),
                  target_output: targets
              })
          seed += 1

        save_path = os.path.join(self.get_temp_dir(),
                                 ("cudnn-rnn-%s-test" % rnn_mode))
        save_v = saver.save(sess, save_path)
        self.assertEqual(save_path, save_v)

    # Cudnn inference graph
    with ops.Graph().as_default() as g:
      model = CudnnTestModel(
          rnn_mode,
          num_layers,
          num_units,
          input_size,
          direction=direction,
          dtype=dtype,
          training=False)
      rnn = model.rnn
      saver = saver_lib.Saver()

      inference_input = np.random.rand(seq_length, batch_size,
                                       input_size).astype(np.float32)
      with self.test_session(use_gpu=True, graph=g) as sess:
        sess.run(variables.global_variables_initializer())
        saver.restore(sess, save_path)

        # Cudnn inference
        cudnn_outputs_v, cudnn_output_states_v = model.Feed(
            sess, inference_input, return_sum=False)

    # Canonical RNN inference graph
    with ops.Graph().as_default() as g:
      cell_inputs = array_ops.placeholder(
          dtype, shape=[seq_length, batch_size, input_size])
      if direction == CUDNN_RNN_UNIDIRECTION:
        # outputs is one tensor, states are num_layer tuples, each 2 tensors
        (outputs, states) = _CreateCudnnCompatibleCanonicalRNN(rnn, cell_inputs)
        if rnn_mode == CUDNN_LSTM:
          output_h = array_ops.stack([s.h for s in states])
          output_c = array_ops.stack([s.c for s in states])
        else:
          output_state = array_ops.stack([s for s in states])
      else:
        # outputs is one tensor.
        # states is a tuple of 2 tuples:
        # each sub tuple is num_layer tuples, each with 2 tensors.
        (outputs, states) = _CreateCudnnCompatibleCanonicalRNN(
            rnn, cell_inputs, is_bidi=True)
        output_state_fw, output_state_bw = states
        if rnn_mode == CUDNN_LSTM:
          output_h, output_c = [], []
          for s_fw, s_bw in zip(output_state_fw, output_state_bw):
            output_h.append(array_ops.stack([s_fw.h, s_bw.h]))
            output_c.append(array_ops.stack([s_fw.c, s_bw.c]))
          output_h = array_ops.concat(output_h, axis=0)
          output_c = array_ops.concat(output_c, axis=0)
        else:
          output_state = []
          for s_fw, s_bw in zip(output_state_fw, output_state_bw):
            output_state.append(array_ops.stack([s_fw, s_bw]))
          output_state = array_ops.concat(output_state, axis=0)
      saver = saver_lib.Saver()

      with self.test_session(use_gpu=True, graph=g) as sess:
        saver.restore(sess, save_path)

        # BlockCell inference
        if rnn_mode == CUDNN_LSTM:
          outputs_v, output_h_v, output_c_v = sess.run(
              [outputs, output_h, output_c],
              feed_dict={cell_inputs: inference_input})
          self.assertAllClose(cudnn_outputs_v, outputs_v)
          cudnn_output_h_v, cudnn_output_c_v = cudnn_output_states_v
          self.assertAllClose(cudnn_output_h_v, output_h_v)
          self.assertAllClose(cudnn_output_c_v, output_c_v)
        else:
          outputs_v, output_state_v = sess.run(
              [outputs, output_state],
              feed_dict={cell_inputs: inference_input})
          self.assertAllClose(cudnn_outputs_v, outputs_v, atol=1e-3, rtol=1e-3)
          (cudnn_output_h_v,) = cudnn_output_states_v
          self.assertAllClose(cudnn_output_h_v, output_state_v, atol=1e-3,
                              rtol=1e-3)


class CudnnRNNTestParamsSize(test_util.TensorFlowTestCase):

  def _TestOpaqueParamsSize(self, rnn_mode, num_layers, num_units, input_size,
                            dtype, direction):
    logging.info("Testing one lstm param size with config: %s", locals())
    model = CudnnTestModel(
        rnn_mode,
        num_layers,
        num_units,
        input_size,
        dtype=dtype,
        direction=direction)
    rnn = model.rnn

    # Min param size estimate = sum(weights.size) + sum(biases.size)
    min_params_size = (
        sum(map(np.prod, rnn.canonical_weight_shapes)) +
        sum(sp[0] for sp in rnn.canonical_bias_shapes))

    opaque_params = rnn.trainable_variables[0]
    with self.test_session(use_gpu=True, graph=ops.get_default_graph()):
      variables.global_variables_initializer().run()
      opaque_params_size_v = opaque_params.eval().size
      self.assertLessEqual(min_params_size, opaque_params_size_v)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testOpaqueParamsSize(self):
    test_configs = [
        [4, 200, 200],
        [4, 200, 300],
        [4, 200, 100],
        [1, 100, 200],
        [2, 200, 100],
        [3, 200, 400],
    ]
    directions = [CUDNN_RNN_UNIDIRECTION, CUDNN_RNN_BIDIRECTION]
    dtype_list = [dtypes.float16, dtypes.float32, dtypes.float64]
    rnns = [CUDNN_LSTM, CUDNN_GRU, CUDNN_RNN_RELU, CUDNN_RNN_TANH]
    for (rnn, config, dtype, direction) in itertools.product(
        rnns, test_configs, dtype_list, directions):
      num_layers, num_units, input_size = config
      with ops.Graph().as_default():
        self._TestOpaqueParamsSize(rnn, num_layers, num_units, input_size,
                                   dtype, direction)


class CudnnRNNTestTraining(test_util.TensorFlowTestCase):

  def setUp(self):
    super(CudnnRNNTestTraining, self).setUp()
    self._reset_rnd_gen_state = os.environ.get("TF_CUDNN_RESET_RND_GEN_STATE",
                                               str(False))
    self._rnn_use_v2 = os.environ.get("TF_CUDNN_RNN_USE_V2", "0")

  def tearDown(self):
    super(CudnnRNNTestTraining, self).tearDown()
    os.environ["TF_CUDNN_RESET_RND_GEN_STATE"] = self._reset_rnd_gen_state
    os.environ["TF_CUDNN_RNN_USE_V2"] = self._rnn_use_v2

  def _ComputeNumericGrad(self, sess, y, x, delta=1e-4, step=1):
    """Compute the numeric gradient of y wrt to x.

    Args:
      sess: The TF session constructed with a graph containing x and y.
      y: A scalar TF Tensor in the graph constructed in sess.
      x: A TF Tensor in the graph constructed in sess.
      delta: Gradient checker's small perturbation of x[i].
      step: Only compute numerical gradients for a subset of x values.
        I.e. dy/dx[i] is computed if i % step == 0.
    Returns:
      A Tensor of the same shape and dtype as x. If x[i] is not chosen
      to compute the numerical gradient dy/x[i], the corresponding
      value is set to 0.
    """

    x_data = sess.run(x)
    x_size = x_data.size
    x_shape = x_data.shape

    numeric_grad = np.zeros(x_size, dtype=x_data.dtype)

    for i in range(0, x_size, step):
      x_pos = x_data.copy()
      if x_size == 1:
        x_pos += delta
      else:
        x_pos.flat[i] += delta
      y_pos_feed_dict = dict([(x.name, x_pos)])
      y_pos = sess.run(y, feed_dict=y_pos_feed_dict)

      x_neg = x_data.copy()
      if x_size == 1:
        x_neg -= delta
      else:
        x_neg.flat[i] -= delta
      y_neg_feed_dict = dict([(x.name, x_neg)])
      y_neg = sess.run(y, feed_dict=y_neg_feed_dict)
      numeric_grad[i] = (y_pos - y_neg) / (2 * delta)
    return numeric_grad.reshape(x_shape)

  def _GetShape(self, sess, inputs):
    if not isinstance(inputs, collections.Iterable):
      return sess.run(array_ops.shape(inputs))
    else:
      return sess.run([array_ops.shape(x) for x in inputs])

  def _GradientCheckFp16(self, sess, y, xs, num_samples,
                         tolerance=1e-6, delta=1e-4):
    """Gradient check for Fp16.

    Fp16 numerical gradients end up being zeros. Use a new way to check
    gradients:

    Given multi-variant function:
    y = f(x1, x2, ... xn)
    delta_y = f(x1 + delta_x1, x2+delta_x2, ..., xn+delta_xn) -
              f(x1, x2, ..., xn)
            = f'(x1) * delta_x1 + f'(x2) * delta_x2 + .. + f'(xn) * delta_xn
    where:
      delta_xi are very small disturbance.
      f'(xi) is the gradient of y w.r.t xi.

    The gradient check verifies the expected delta_y calculated by the above
    equation is close to the actual delta_y.
    Args:
      sess: tf.Session object.
      y: output tensor.
      xs: a tensor or a list of input tensors.
      num_samples: number of test samples to run.
      tolerance: error tolerance.
      delta: the order of magnititued of input disturbance to apply to calculate
        the output change w.r.t inputs.
    """
    sym_grads = self._ComputeSymGrads(sess, y, xs)
    xs_shapes = self._GetShape(sess, xs)

    x_vals = [sess.run(x) for x in xs]
    for _ in range(num_samples):
      delta_xs = [delta * np.random.rand(*shape.tolist())
                  for shape in xs_shapes]

      feed_dict = {}
      for x, x_val, delta_x in zip(xs, x_vals, delta_xs):
        feed_dict[x] = x_val + delta_x
      actual_delta_y = (float(sess.run(y, feed_dict=feed_dict)) -
                        float(sess.run(y)))

      expected_delta_y = 0.
      for sym_grad, delta_x in zip(sym_grads, delta_xs):
        expected_delta_y += np.dot(
            sym_grad.astype(np.float32).flatten(),
            delta_x.astype(np.float32).flatten())
      self.assertAllClose(expected_delta_y, actual_delta_y,
                          atol=tolerance, rtol=tolerance)

  def _GradientCheck(self, sess, y, xs, tolerance=1e-6, delta=1e-4):
    sym_grads = self._ComputeSymGrads(sess, y, xs)

    num_grads = [self._ComputeNumericGrad(sess, y, x, delta) for x in xs]
    self.assertEqual(len(sym_grads), len(num_grads))
    for x, sym, num in zip(xs, sym_grads, num_grads):
      logging.info("Comparing gradients for input: %s", x.name)
      self.assertFalse(np.any(np.isnan(sym)))
      self.assertFalse(np.any(np.isnan(num)))
      self.assertAllClose(sym, num, atol=tolerance, rtol=tolerance)

  def _ComputeSymGrads(self, sess, y, xs):
    sym_grads_t = gradients.gradients(y, xs)
    return sess.run(sym_grads_t)

  def _TestOneSimpleTraining(self, rnn_mode, num_layers, num_units, input_size,
                             batch_size, seq_length, dir_count, dropout, dtype,
                             use_v2, delta, tolerance):
    # Gradient checking runs two forward ops with almost the same input. Need to
    # make sure the drop patterns across the two runs are the same.
    logging.info("Training test with config: %s", locals())
    os.environ["TF_CUDNN_RESET_RND_GEN_STATE"] = str(True)

    np.random.seed(1234)
    random_seed.set_random_seed(5678)
    has_input_c = (rnn_mode == CUDNN_LSTM)
    direction = (CUDNN_RNN_UNIDIRECTION
                 if dir_count == 1 else CUDNN_RNN_BIDIRECTION)
    if use_v2:
      os.environ["TF_CUDNN_RNN_USE_V2"] = "1"
    else:
      os.environ["TF_CUDNN_RNN_USE_V2"] = "0"
    model = CudnnTestModel(
        rnn_mode,
        num_layers,
        num_units,
        input_size,
        direction=direction,
        dropout=dropout,
        dtype=dtype,
        training=True,
        bias_initializer=init_ops.random_normal_initializer(
            mean=1., dtype=dtype))
    rnn = model.rnn
    params = rnn.trainable_variables[0]

    inputs = variables.Variable(
        random_ops.random_uniform([seq_length, batch_size, input_size],
                                  dtype=dtype),
        dtype=dtype).read_value()
    input_h = variables.Variable(
        random_ops.random_uniform(
            [num_layers * dir_count, batch_size, num_units], dtype=dtype),
        dtype=dtype).read_value()
    if has_input_c:
      input_c = variables.Variable(
          random_ops.random_uniform(
              [num_layers * dir_count, batch_size, num_units], dtype=dtype),
          dtype=dtype).read_value()
      initial_state = (input_h, input_c)
    else:
      initial_state = (input_h,)
    total_sum = model.FProp(inputs, initial_state, training=True)

    with self.test_session(use_gpu=True, graph=ops.get_default_graph()) as sess:
      sess.run(variables.global_variables_initializer())
      all_inputs = [inputs, params]
      for s in initial_state:
        all_inputs.append(s)
      if dtype == dtypes.float16:
        self._GradientCheckFp16(
            sess, total_sum, all_inputs,
            num_samples=FLAGS.grad_check_num_samples,
            tolerance=tolerance, delta=delta)
      else:
        for _ in range(FLAGS.grad_check_num_samples):
          # Each time choose a different set of inputs.
          sess.run(variables.global_variables_initializer())
          self._GradientCheck(
              sess, total_sum, all_inputs,
              tolerance=tolerance, delta=delta)

  def _TestSimpleTrainingHelper(self, rnn_mode, test_configs):
    dropouts = [0, 0.5, 1.]
    v2_options = [False, True]
    for config, dropout, use_v2 in itertools.product(test_configs, dropouts,
                                                     v2_options):
      dtype = config.get("dtype", dtypes.float32)
      delta = config.get("delta", 1e-4)
      tolerance = config.get("tolerance", 1e-6)
      dir_count = config.get("dir_count", 1)
      shape = config["shape"]
      if dtype == dtypes.float64:
        # TODO(jamesqin): b/117848763
        use_v2 = False
      with ops.Graph().as_default():
        self._TestOneSimpleTraining(
            rnn_mode, shape["num_layers"], shape["num_units"],
            shape["input_size"], shape["batch_size"], shape["seq_length"],
            dir_count, dropout, dtype, use_v2, delta, tolerance)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testSimpleTrainingLSTMFp64(self):
    test_configs = [
        {
            "dtype": dtypes.float64,
            "tolerance": 5e-6,
            "shape": {
                "num_layers": 2,
                "num_units": 3,
                "input_size": 4,
                "batch_size": 3,
                "seq_length": 4,
            },
        },
    ]
    self._TestSimpleTrainingHelper(CUDNN_LSTM, test_configs)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testSimpleTrainingLSTMFp32(self):
    test_configs = [
        {
            "dtype": dtypes.float32,
            "delta": 1e-4,
            "tolerance": 9e-2,
            "shape": {
                "num_layers": 2,
                "num_units": 3,
                "input_size": 4,
                "batch_size": 3,
                "seq_length": 4,
            },
        },
    ]
    self._TestSimpleTrainingHelper(CUDNN_LSTM, test_configs)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testSimpleTrainingLSTMFp16(self):
    test_configs = [
        {
            "dtype": dtypes.float16,
            "delta": 1e-3,
            "tolerance": 9e-2,
            "shape": {
                "num_layers": 2,
                "num_units": 3,
                "input_size": 4,
                "batch_size": 3,
                "seq_length": 4,
            },
        },
        {
            "dtype": dtypes.float16,
            "delta": 1e-2,
            "tolerance": 9e-2,
            "shape": {
                "num_layers": 2,
                "num_units": 6,
                "input_size": 8,
                "batch_size": 6,
                "seq_length": 4,
            },
        },
    ]
    self._TestSimpleTrainingHelper(CUDNN_LSTM, test_configs)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testSimpleTrainingGRUFp64(self):
    test_configs = [
        {
            "dtype": dtypes.float64,
            "tolerance": 5e-6,
            "shape": {
                "num_layers": 2,
                "num_units": 3,
                "input_size": 4,
                "batch_size": 3,
                "seq_length": 4,
            }
        },
    ]
    self._TestSimpleTrainingHelper(CUDNN_GRU, test_configs)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testSimpleTrainingGRUFp32(self):
    test_configs = [
        {
            "dtype": dtypes.float32,
            "delta": 1e-3,
            "tolerance": 4e-3,
            "shape": {
                "num_layers": 2,
                "num_units": 3,
                "input_size": 4,
                "batch_size": 3,
                "seq_length": 4,
            },
        },
    ]
    self._TestSimpleTrainingHelper(CUDNN_GRU, test_configs)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testSimpleTrainingGRUFp16(self):
    test_configs = [
        {
            "dtype": dtypes.float16,
            "delta": 2e-3,
            "tolerance": 6e-2,
            "shape": {
                "num_layers": 2,
                "num_units": 3,
                "input_size": 4,
                "batch_size": 3,
                "seq_length": 4,
            },
        },
    ]
    self._TestSimpleTrainingHelper(CUDNN_GRU, test_configs)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testSimpleTrainingRNNTanhFp64(self):
    test_configs = [
        {
            "dtype": dtypes.float64,
            "tolerance": 5e-6,
            "shape": {
                "num_layers": 2,
                "num_units": 3,
                "input_size": 4,
                "batch_size": 3,
                "seq_length": 4,
            },
        },
    ]
    self._TestSimpleTrainingHelper(CUDNN_RNN_TANH, test_configs)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testSimpleTrainingRNNTanhFp32(self):
    test_configs = [
        {
            "dtype": dtypes.float32,
            "delta": 1e-3,
            "tolerance": 5e-3,
            "shape": {
                "num_layers": 2,
                "num_units": 3,
                "input_size": 4,
                "batch_size": 3,
                "seq_length": 4,
            },
        },
    ]
    self._TestSimpleTrainingHelper(CUDNN_RNN_TANH, test_configs)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testSimpleTrainingRNNTanhFp16(self):
    test_configs = [
        {
            "dtype": dtypes.float16,
            "delta": 1e-3,
            "tolerance": 5e-2,
            "shape": {
                "num_layers": 2,
                "num_units": 3,
                "input_size": 4,
                "batch_size": 3,
                "seq_length": 4,
            },
        },
    ]
    self._TestSimpleTrainingHelper(CUDNN_RNN_TANH, test_configs)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testSimpleTrainingRNNReluFp64(self):
    test_configs = [
        {
            "dtype": dtypes.float64,
            "tolerance": 5e-6,
            "shape": {
                "num_layers": 2,
                "num_units": 3,
                "input_size": 4,
                "batch_size": 3,
                "seq_length": 4,
            },
        },
    ]
    self._TestSimpleTrainingHelper(CUDNN_RNN_RELU, test_configs)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testSimpleTrainingRNNReluFp32(self):
    test_configs = [
        {
            "dtype": dtypes.float32,
            "delta": 1e-4,
            "tolerance": 3e-1,
            "shape": {
                "num_layers": 2,
                "num_units": 3,
                "input_size": 4,
                "batch_size": 3,
                "seq_length": 4,
            },
        },
    ]
    self._TestSimpleTrainingHelper(CUDNN_RNN_RELU, test_configs)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testSimpleTrainingRNNReluFp16(self):
    test_configs = [
        {
            "dtype": dtypes.float16,
            "delta": 1e-3,
            "tolerance": 7e-2,
            "shape": {
                "num_layers": 2,
                "num_units": 3,
                "input_size": 4,
                "batch_size": 3,
                "seq_length": 4,
            },
        },
    ]
    self._TestSimpleTrainingHelper(CUDNN_RNN_RELU, test_configs)


if __name__ == "__main__":
  argv0 = sys.argv[0]
  parser = argparse.ArgumentParser()
  parser.add_argument(
      "--grad_check_num_samples",
      type=int,
      default=1,
      help="Number of samples to run for gradient check.")
  FLAGS, unparsed = parser.parse_known_args()
  sys.argv = [argv0] + unparsed
  googletest.main()
