# 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 itertools
import os
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.framework import dtypes
from tensorflow.python.framework import ops
from tensorflow.python.framework import random_seed
from tensorflow.python.framework.test_util import TensorFlowTestCase
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 gradient_descent
from tensorflow.python.training import saver as saver_lib

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,
               kernel_initializer=None,
               bias_initializer=None):
    if dtype not in (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,
        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(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)


# TODO(jamesqin): Transform to parameterized test after it is included in the
# TF open source codebase.
class CudnnRNNTestSaveRestore(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._OpaqueParamsToCanonical()
      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 explictly, 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])
      weights1, biases1 = model1.rnn.saveable._OpaqueParamsToCanonical()
      weights2, biases2 = model2.rnn.saveable._OpaqueParamsToCanonical()
      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 explictly, 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 explictly, 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.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)


# TODO(jamesqin): Transform to parameterized test after it is included in the
# TF open source codebase.
class CudnnRNNTestCompatibleRNNCells(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-5, rtol=1e-5)
          (cudnn_output_h_v,) = cudnn_output_states_v
          self.assertAllClose(cudnn_output_h_v, output_state_v, atol=1e-5,
                              rtol=1e-5)


class CudnnRNNTestParamsSize(TensorFlowTestCase):

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

    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 = (
        np.sum(map(np.prod, rnn.canonical_weight_shapes)) +
        np.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]
    rnns = [CUDNN_LSTM, CUDNN_GRU, CUDNN_RNN_RELU, CUDNN_RNN_TANH]
    for (rnn, config, direction) in itertools.product(rnns, test_configs,
                                                      directions):
      num_layers, num_units, input_size = config
      with ops.Graph().as_default():
        self._TestOpaqueParamsSize(rnn, num_layers, num_units, input_size,
                                   direction)


class CudnnRNNTestTraining(TensorFlowTestCase):

  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 _GradientCheck(self, sess, y, xs, tolerance=1e-6, delta=1e-4):
    sym_grads_t = gradients.gradients(y, xs)
    sym_grads = sess.run(sym_grads_t)

    num_grads = [self._ComputeNumericGrad(sess, y, x, delta) for x in xs]
    self.assertEqual(len(sym_grads), len(num_grads))
    for sym, num in zip(sym_grads, num_grads):
      self.assertFalse(np.any(np.isnan(sym)))
      self.assertFalse(np.any(np.isnan(num)))
      self.assertAllClose(sym, num, atol=tolerance, rtol=tolerance)

  def _TestOneSimpleTraining(self, rnn_mode, num_layers, num_units, input_size,
                             batch_size, seq_length, dir_count, dropout, dtype,
                             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())
    old_env_state = os.environ.get("TF_CUDNN_RESET_RND_GEN_STATE", str(False))
    os.environ["TF_CUDNN_RESET_RND_GEN_STATE"] = str(True)
    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)
    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)
    input_h = variables.Variable(
        random_ops.random_uniform(
            [num_layers * dir_count, batch_size, num_units], dtype=dtype),
        dtype=dtype)
    if has_input_c:
      input_c = variables.Variable(
          random_ops.random_uniform(
              [num_layers * dir_count, batch_size, num_units], dtype=dtype),
          dtype=dtype)
      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)
      self._GradientCheck(
          sess, total_sum, all_inputs, tolerance=tolerance, delta=delta)
      os.environ["TF_CUDNN_RESET_RND_GEN_STATE"] = old_env_state

  def _TestSimpleTrainingHelper(self, rnn_mode, test_configs):
    dropouts = [0., 0.5, 1.]
    for config, dropout in itertools.product(test_configs, dropouts):
      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"]
      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, delta, tolerance)

  @unittest.skipUnless(test.is_built_with_cuda(),
                       "Test only applicable when running on GPUs")
  def testSimpleTrainingLSTM64(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 testSimpleTrainingLSTM32(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 testSimpleTrainingGRU64(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 testSimpleTrainingGRU32(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 testSimpleTrainingRNNTanh64(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 testSimpleTrainingRNNTanh32(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 testSimpleTrainingRNNRelu64(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 testSimpleTrainingRNNRelu32(self):
    test_configs = [
        {
            "dtype": dtypes.float32,
            "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__":
  googletest.main()
