# Copyright 2018 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 GRU V2 layer."""

from __future__ import absolute_import
from __future__ import division
from __future__ import print_function

import os
import shutil

from absl.testing import parameterized
import numpy as np

from tensorflow.core.protobuf import config_pb2
from tensorflow.core.protobuf import rewriter_config_pb2
from tensorflow.python import keras
from tensorflow.python.data.ops import dataset_ops
from tensorflow.python.eager import backprop
from tensorflow.python.eager import context
from tensorflow.python.framework import constant_op
from tensorflow.python.framework import dtypes
from tensorflow.python.framework import random_seed
from tensorflow.python.framework import test_util
from tensorflow.python.keras import keras_parameterized
from tensorflow.python.keras import testing_utils
from tensorflow.python.keras.layers import recurrent as rnn_v1
from tensorflow.python.keras.layers import recurrent_v2 as rnn
from tensorflow.python.ops import array_ops
from tensorflow.python.ops import control_flow_ops
from tensorflow.python.ops import gen_math_ops
from tensorflow.python.ops import math_ops
from tensorflow.python.ops import random_ops
from tensorflow.python.platform import test
from tensorflow.python.training import gradient_descent


# Global config for grappler setting that is used for graph mode test.
_rewrites = rewriter_config_pb2.RewriterConfig()
_rewrites.implementation_selector = rewriter_config_pb2.RewriterConfig.ON
_rewrites.min_graph_nodes = -1
_graph_options = config_pb2.GraphOptions(rewrite_options=_rewrites)
_config = config_pb2.ConfigProto(graph_options=_graph_options)


@keras_parameterized.run_all_keras_modes(config=_config)
class GRUV2Test(keras_parameterized.TestCase):

  @parameterized.named_parameters(
      ('non_tan_activation', 'relu', 'sigmoid', 0, False, True, True),
      ('non_sigmoid_recur_activation', 'tanh', 'relu', 0, False, True, True),
      ('use_recurrent_dropout', 'tanh', 'sigmoid', 0.1, False, True, True),
      ('unroll', 'tanh', 'sigmoid', 0, True, True, True),
      ('not_use_bias', 'tanh', 'sigmoid', 0, False, False, True),
      ('not_reset_after', 'tanh', 'sigmoid', 0, False, True, False)
  )
  def test_could_use_defun_backend(self, activation, recurrent_activation,
                                   recurrent_dropout, unroll, use_bias,
                                   reset_after):
    layer = rnn.GRU(1,
                    activation=activation,
                    recurrent_activation=recurrent_activation,
                    recurrent_dropout=recurrent_dropout,
                    unroll=unroll,
                    use_bias=use_bias,
                    reset_after=reset_after)
    self.assertFalse(layer.could_use_cudnn)

  def test_keras_model_with_gru(self):
    input_shape = 10
    rnn_state_size = 8
    output_shape = 8
    timestep = 4
    batch = 100
    epoch = 10

    (x_train, y_train), _ = testing_utils.get_test_data(
        train_samples=batch,
        test_samples=0,
        input_shape=(timestep, input_shape),
        num_classes=output_shape)
    y_train = keras.utils.to_categorical(y_train, output_shape)

    layer = rnn.GRU(rnn_state_size)

    inputs = keras.layers.Input(
        shape=[timestep, input_shape], dtype=dtypes.float32)

    outputs = layer(inputs)
    model = keras.models.Model(inputs, outputs)
    model.compile('rmsprop', loss='mse')
    model.fit(x_train, y_train, epochs=epoch)
    model.evaluate(x_train, y_train)
    model.predict(x_train)

  def test_dynamic_behavior_GRU(self):
    num_samples = 2
    timesteps = 3
    embedding_dim = 4
    units = 2
    layer = rnn.GRU(units, input_shape=(None, embedding_dim))
    model = keras.models.Sequential()
    model.add(layer)
    model.compile(gradient_descent.GradientDescentOptimizer(0.001), 'mse')
    x = np.random.random((num_samples, timesteps, embedding_dim))
    y = np.random.random((num_samples, units))
    model.train_on_batch(x, y)

  def test_stacking_GRU(self):
    inputs = np.random.random((2, 3, 4))
    targets = np.abs(np.random.random((2, 3, 5)))
    targets /= targets.sum(axis=-1, keepdims=True)
    model = keras.models.Sequential()
    model.add(rnn.GRU(10, return_sequences=True, unroll=False))
    model.add(rnn.GRU(5, return_sequences=True, unroll=False))
    model.compile(
        loss='categorical_crossentropy',
        optimizer=gradient_descent.GradientDescentOptimizer(0.01))
    model.fit(inputs, targets, epochs=1, batch_size=2, verbose=1)

  def test_from_config_GRU(self):
    layer_class = rnn.GRU
    for stateful in (False, True):
      l1 = layer_class(units=1, stateful=stateful)
      l2 = layer_class.from_config(l1.get_config())
      assert l1.get_config() == l2.get_config()

  @test_util.run_v2_only
  def test_gru_v2_feature_parity_with_canonical_gru(self):
    input_shape = 10
    rnn_state_size = 8
    timestep = 4
    batch = 20

    (x_train, y_train), _ = testing_utils.get_test_data(
        train_samples=batch,
        test_samples=0,
        input_shape=(timestep, input_shape),
        num_classes=rnn_state_size,
        random_seed=random_seed.DEFAULT_GRAPH_SEED)
    y_train = keras.utils.to_categorical(y_train, rnn_state_size)
    # For the last batch item of the test data, we filter out the last
    # timestep to simulate the variable length sequence and masking test.
    x_train[-2:, -1, :] = 0.0
    y_train[-2:] = 0

    inputs = keras.layers.Input(
        shape=[timestep, input_shape], dtype=dtypes.float32)
    masked_input = keras.layers.Masking()(inputs)
    gru_layer = rnn_v1.GRU(rnn_state_size,
                           recurrent_activation='sigmoid',
                           reset_after=True)
    output = gru_layer(masked_input)
    gru_model = keras.models.Model(inputs, output)
    weights = gru_model.get_weights()
    y_1 = gru_model.predict(x_train)
    gru_model.compile('rmsprop', 'mse')
    gru_model.fit(x_train, y_train)
    y_2 = gru_model.predict(x_train)

    with test_util.device(use_gpu=True):
      cudnn_layer = rnn.GRU(rnn_state_size,
                            recurrent_activation='sigmoid',
                            reset_after=True)
      cudnn_model = keras.models.Model(inputs, cudnn_layer(masked_input))
    cudnn_model.set_weights(weights)
    y_3 = cudnn_model.predict(x_train)
    cudnn_model.compile('rmsprop', 'mse')
    cudnn_model.fit(x_train, y_train)
    y_4 = cudnn_model.predict(x_train)

    self.assertAllClose(y_1, y_3, rtol=2e-5, atol=2e-5)
    self.assertAllClose(y_2, y_4, rtol=2e-5, atol=2e-5)

  @parameterized.named_parameters(
      # test_name, use_bias, bias_initializer, activation
      ('normal', True, 'zeros'),
      ('no_bias', False, 'zeros'),
      ('random_bias', True, 'random_uniform'),
  )
  def test_gru_v2_model_save_load(self, use_bias, bias_initializer):
    temp_dir = self.get_temp_dir()
    self.addCleanup(shutil.rmtree, temp_dir)
    h5_path = os.path.join(temp_dir, 'test.h5')

    batch = 10
    timestep = 3
    input_dim = 5
    units = 2

    x = np.random.random((batch, timestep, input_dim))

    def build_model():
      inputs = keras.layers.Input(
          shape=[timestep, input_dim], dtype=dtypes.float32)
      layer = rnn.GRU(
          units,
          use_bias=use_bias,
          bias_initializer=bias_initializer)
      output = layer(inputs)
      return keras.models.Model(inputs, output), layer

    model, layer = build_model()
    y_ref = model.predict(x)
    model.save_weights(h5_path)

    cloned_model, new_layer = build_model()
    cloned_model.load_weights(h5_path)
    y = cloned_model.predict(x)

    self.assertAllClose(y, y_ref)
    self.assertAllClose(layer.get_weights(), new_layer.get_weights())

  def test_gru_v2_output_on_multiple_kernel(self):
    input_shape = 10
    rnn_state_size = 8
    timestep = 4
    batch = 100

    x_train = np.random.random((batch, timestep, input_shape))

    inputs = keras.layers.Input(
        shape=[timestep, input_shape], dtype=dtypes.float32)
    with test_util.device(use_gpu=False):
      layer = rnn.GRU(rnn_state_size)
      output = layer(inputs)
      cpu_model = keras.models.Model(inputs, output)
      weights = cpu_model.get_weights()
      y_1 = cpu_model.predict(x_train)

    with test_util.device(use_gpu=True):
      layer = rnn.GRU(rnn_state_size)
      output = layer(inputs)
      gpu_model = keras.models.Model(inputs, output)
      gpu_model.set_weights(weights)
      y_2 = gpu_model.predict(x_train)

    # Note that CuDNN uses 'sigmoid' as activation, so the GRU V2 uses
    # 'sigmoid' as default. Construct the canonical GRU with sigmoid to achieve
    # the same output.
    with test_util.device(use_gpu=True):
      layer = rnn_v1.GRU(rnn_state_size,
                         recurrent_activation='sigmoid',
                         reset_after=True)
      output = layer(inputs)
      canonical_model = keras.models.Model(inputs, output)
      canonical_model.set_weights(weights)
      y_3 = canonical_model.predict(x_train)

    self.assertAllClose(y_1, y_2, rtol=1e-5, atol=1e-5)
    self.assertAllClose(y_2, y_3, rtol=1e-5, atol=1e-5)

  @parameterized.named_parameters(
      # test_name, time_major, go_backwards
      ('normal', False, False),
      ('time_major', True, False),
      ('go_backwards', False, True),
      ('both', True, True),
  )
  def test_time_major_and_go_backward(self, time_major, go_backwards):
    input_shape = 10
    rnn_state_size = 8
    timestep = 4
    batch = 100

    x_train = np.random.random((batch, timestep, input_shape))

    def build_model(layer_cls):
      inputs = keras.layers.Input(
          shape=[timestep, input_shape], dtype=dtypes.float32)
      layer = layer_cls(rnn_state_size,
                        recurrent_activation='sigmoid',
                        time_major=time_major,
                        return_sequences=True,
                        go_backwards=go_backwards,
                        reset_after=True)
      if time_major:
        converted_input = keras.layers.Lambda(
            lambda t: array_ops.transpose(t, [1, 0, 2]))(inputs)
        outputs = layer(converted_input)
        outputs = keras.layers.Lambda(
            lambda t: array_ops.transpose(t, [1, 0, 2]))(outputs)
      else:
        outputs = layer(inputs)
      return keras.models.Model(inputs, outputs)

    gru_model = build_model(rnn_v1.GRU)
    y_ref = gru_model.predict(x_train)
    weights = gru_model.get_weights()

    gru_v2_model = build_model(rnn.GRU)
    gru_v2_model.set_weights(weights)
    y = gru_v2_model.predict(x_train)

    self.assertAllClose(y, y_ref)

  def test_with_masking_layer_GRU(self):
    layer_class = rnn.GRU
    inputs = np.random.random((2, 3, 4))
    targets = np.abs(np.random.random((2, 3, 5)))
    targets /= targets.sum(axis=-1, keepdims=True)
    model = keras.models.Sequential()
    model.add(keras.layers.Masking(input_shape=(3, 4)))
    model.add(layer_class(units=5, return_sequences=True, unroll=False))
    model.compile(loss='categorical_crossentropy',
                  optimizer=gradient_descent.GradientDescentOptimizer(0.001))
    model.fit(inputs, targets, epochs=1, batch_size=2, verbose=1)

  def test_masking_with_stacking_GRU(self):
    inputs = np.random.random((2, 3, 4))
    targets = np.abs(np.random.random((2, 3, 5)))
    targets /= targets.sum(axis=-1, keepdims=True)
    model = keras.models.Sequential()
    model.add(keras.layers.Masking(input_shape=(3, 4)))
    model.add(rnn.GRU(10, return_sequences=True, unroll=False))
    model.add(rnn.GRU(5, return_sequences=True, unroll=False))
    model.compile(
        loss='categorical_crossentropy',
        optimizer=gradient_descent.GradientDescentOptimizer(0.01))
    model.fit(inputs, targets, epochs=1, batch_size=2, verbose=1)

  def test_return_sequences_GRU(self):
    num_samples = 2
    timesteps = 3
    embedding_dim = 4
    units = 2
    testing_utils.layer_test(
        rnn.GRU,
        kwargs={'units': units,
                'return_sequences': True},
        input_shape=(num_samples, timesteps, embedding_dim))

  def test_float64_GRU(self):
    num_samples = 2
    timesteps = 3
    embedding_dim = 4
    units = 2
    testing_utils.layer_test(
        rnn.GRU,
        kwargs={'units': units,
                'return_sequences': True,
                'dtype': 'float64'},
        input_shape=(num_samples, timesteps, embedding_dim),
        input_dtype='float64')

  def test_return_states_GRU(self):
    layer_class = rnn.GRU
    x = np.random.random((2, 3, 4))
    y = np.abs(np.random.random((2, 5)))
    s = np.abs(np.random.random((2, 5)))
    inputs = keras.layers.Input(
        shape=[3, 4], dtype=dtypes.float32)
    masked = keras.layers.Masking()(inputs)
    outputs, states = layer_class(units=5, return_state=True)(masked)

    model = keras.models.Model(inputs, [outputs, states])
    model.compile(loss='categorical_crossentropy',
                  optimizer=gradient_descent.GradientDescentOptimizer(0.001))
    model.fit(x, [y, s], epochs=1, batch_size=2, verbose=1)

  def test_dropout_GRU(self):
    num_samples = 2
    timesteps = 3
    embedding_dim = 4
    units = 2
    testing_utils.layer_test(
        rnn.GRU,
        kwargs={'units': units,
                'dropout': 0.1,
                'recurrent_dropout': 0.1},
        input_shape=(num_samples, timesteps, embedding_dim))

  def test_constraints_GRU(self):
    embedding_dim = 4
    layer_class = rnn.GRU
    k_constraint = keras.constraints.max_norm(0.01)
    r_constraint = keras.constraints.max_norm(0.01)
    b_constraint = keras.constraints.max_norm(0.01)
    layer = layer_class(
        5,
        return_sequences=False,
        weights=None,
        input_shape=(None, embedding_dim),
        kernel_constraint=k_constraint,
        recurrent_constraint=r_constraint,
        bias_constraint=b_constraint)
    layer.build((None, None, embedding_dim))
    self.assertEqual(layer.cell.kernel.constraint, k_constraint)
    self.assertEqual(layer.cell.recurrent_kernel.constraint, r_constraint)
    self.assertEqual(layer.cell.bias.constraint, b_constraint)

  @parameterized.parameters([0, 1, 2])
  def test_implementation_mode_GRU(self, implementation_mode):
    num_samples = 2
    timesteps = 3
    embedding_dim = 4
    units = 2
    testing_utils.layer_test(
        rnn.GRU,
        kwargs={'units': units,
                'implementation': implementation_mode},
        input_shape=(num_samples, timesteps, embedding_dim))

  def test_regularizers_GRU(self):
    embedding_dim = 4
    layer_class = rnn.GRU
    layer = layer_class(
        5,
        return_sequences=False,
        weights=None,
        input_shape=(None, embedding_dim),
        kernel_regularizer=keras.regularizers.l1(0.01),
        recurrent_regularizer=keras.regularizers.l1(0.01),
        bias_regularizer='l2',
        activity_regularizer='l1')
    layer.build((None, None, 2))
    self.assertEqual(len(layer.losses), 3)

    x = keras.backend.variable(np.ones((2, 3, 2)))
    layer(x)
    if context.executing_eagerly():
      self.assertEqual(len(layer.losses), 4)
    else:
      self.assertEqual(len(layer.get_losses_for(x)), 1)

  def test_statefulness_GRU(self):
    num_samples = 2
    timesteps = 3
    embedding_dim = 4
    units = 2
    layer_class = rnn.GRU
    model = keras.models.Sequential()
    model.add(
        keras.layers.Embedding(
            4,
            embedding_dim,
            mask_zero=True,
            input_length=timesteps,
            batch_input_shape=(num_samples, timesteps)))
    layer = layer_class(
        units, return_sequences=False, stateful=True, weights=None)
    model.add(layer)
    model.compile(
        optimizer=gradient_descent.GradientDescentOptimizer(0.01),
        loss='mse',
        run_eagerly=testing_utils.should_run_eagerly(),
        experimental_run_tf_function=testing_utils.should_run_tf_function())
    out1 = model.predict(np.ones((num_samples, timesteps)))
    self.assertEqual(out1.shape, (num_samples, units))

    # train once so that the states change
    model.train_on_batch(
        np.ones((num_samples, timesteps)), np.ones((num_samples, units)))
    out2 = model.predict(np.ones((num_samples, timesteps)))

    # if the state is not reset, output should be different
    self.assertNotEqual(out1.max(), out2.max())

    # check that output changes after states are reset
    # (even though the model itself didn't change)
    layer.reset_states()
    out3 = model.predict(np.ones((num_samples, timesteps)))
    self.assertNotEqual(out2.max(), out3.max())

    # check that container-level reset_states() works
    model.reset_states()
    out4 = model.predict(np.ones((num_samples, timesteps)))
    np.testing.assert_allclose(out3, out4, atol=1e-5)

    # check that the call to `predict` updated the states
    out5 = model.predict(np.ones((num_samples, timesteps)))
    self.assertNotEqual(out4.max(), out5.max())

    # Check masking
    layer.reset_states()

    left_padded_input = np.ones((num_samples, timesteps))
    left_padded_input[0, :1] = 0
    left_padded_input[1, :2] = 0
    out6 = model.predict(left_padded_input)

    layer.reset_states()

    right_padded_input = np.ones((num_samples, timesteps))
    right_padded_input[0, -1:] = 0
    right_padded_input[1, -2:] = 0
    out7 = model.predict(right_padded_input)

    layer.reset_states()

    mix_padded_input = np.ones((num_samples, timesteps))
    mix_padded_input[0, 1] = 0
    mix_padded_input[1, 0] = 0
    mix_padded_input[1, 2] = 0
    out8 = model.predict(mix_padded_input)

    self.assertAllClose(out7, out6, atol=1e-5)
    self.assertAllClose(out8, out7, atol=1e-5)

  def test_stateful_GRU_training(self):
    # See b/123587692 for more context.
    vocab_size = 20
    embedding_dim = 10
    batch_size = 8
    timestep = 12
    units = 5
    x = np.random.randint(0, vocab_size, size=(batch_size, timestep))
    y = np.random.randint(0, vocab_size, size=(batch_size, timestep))

    model = keras.Sequential([
        keras.layers.Embedding(vocab_size, embedding_dim,
                               batch_input_shape=[batch_size, timestep]),
        rnn.GRU(units, return_sequences=True, stateful=True),
        keras.layers.Dense(vocab_size)
    ])
    model.compile(
        optimizer='adam',
        loss='sparse_categorical_crossentropy',
        run_eagerly=testing_utils.should_run_eagerly(),
        experimental_run_tf_function=testing_utils.should_run_tf_function())
    model.fit(x, y, epochs=1, shuffle=False)

  @test_util.run_v2_only
  def test_explicit_device_with_go_backward_and_mask(self):
    batch_size = 8
    timestep = 7
    masksteps = 5
    units = 4

    inputs = np.random.randn(batch_size, timestep, units).astype(np.float32)
    mask = np.ones((batch_size, timestep)).astype(np.bool)
    mask[:, masksteps:] = 0

    # Test for V1 behavior.
    lstm_v1 = rnn_v1.GRU(units, return_sequences=True, go_backwards=True)
    with test_util.device(use_gpu=True):
      outputs_masked_v1 = lstm_v1(inputs, mask=constant_op.constant(mask))
      outputs_trimmed_v1 = lstm_v1(inputs[:, :masksteps])
    self.assertAllClose(outputs_masked_v1[:, -masksteps:], outputs_trimmed_v1)

    # Test for V2 behavior.
    lstm = rnn.GRU(units, return_sequences=True, go_backwards=True)
    with test_util.device(use_gpu=True):
      outputs_masked = lstm(inputs, mask=constant_op.constant(mask))
      outputs_trimmed = lstm(inputs[:, :masksteps])
    self.assertAllClose(outputs_masked[:, -masksteps:], outputs_trimmed)

  @test_util.run_deprecated_v1
  def test_v1_session_behavior(self):
    # See b/139132348 for more details.
    x = np.random.uniform(size=(100, 4, 8))
    y = np.random.uniform(size=(100, 1))
    dataset = dataset_ops.Dataset.from_tensor_slices(
        (x, y)).shuffle(100).batch(32)

    inp = keras.layers.Input(shape=(4, 8))
    layer = rnn.GRU(1)(inp)
    layer = keras.layers.Dense(1)(layer)

    model = keras.models.Model(inp, layer)

    model.compile(loss='mse', optimizer='sgd')
    model.fit(dataset)


class GRULayerGradientTapeTest(test.TestCase):

  @test_util.run_in_graph_and_eager_modes(config=_config)
  def test_in_tape(self):
    if not context.executing_eagerly():
      self.skipTest('bloo')
    time_steps = 10
    embedding_size = 11
    gru_unit_size = 12

    gru = rnn.GRU(gru_unit_size,
                  return_sequences=True,
                  return_state=True,
                  recurrent_activation='sigmoid',
                  recurrent_initializer='glorot_uniform')

    x = random_ops.random_uniform([1, time_steps, embedding_size])
    y = random_ops.random_uniform([1, gru_unit_size])

    with backprop.GradientTape() as tape:
      hidden_state = array_ops.zeros([1, gru_unit_size], dtype=dtypes.float32)
      _, state = gru(x, initial_state=hidden_state)

      loss = math_ops.reduce_mean(math_ops.square(state - y))

    tape.gradient(loss, gru.variables)


@keras_parameterized.run_all_keras_modes(config=_config)
class GRUGraphRewriteTest(keras_parameterized.TestCase):

  input_shape = 10
  output_shape = 8
  rnn_state_size = 8
  timestep = 4
  batch = 100
  epoch = 1

  def _test_runtime_with_model(self, model):
    (x_train, y_train), _ = testing_utils.get_test_data(
        train_samples=self.batch,
        test_samples=0,
        input_shape=(self.timestep, self.input_shape),
        num_classes=self.output_shape)
    y_train = keras.utils.to_categorical(y_train, self.output_shape)

    model.compile(
        optimizer='sgd',
        loss=['categorical_crossentropy', None],
        experimental_run_tf_function=testing_utils.should_run_tf_function())

    existing_loss = 0
    for _ in range(self.epoch):
      history = model.fit(x_train, y_train)
      loss_value = history.history['loss'][0]

      self.assertNotEqual(existing_loss, loss_value)
      existing_loss = loss_value

    _, runtime_value = model.predict(x_train)
    if test.is_gpu_available():
      self.assertEqual(runtime_value[0], rnn._RUNTIME_GPU)
    else:
      self.assertEqual(runtime_value[0], rnn._RUNTIME_CPU)

  @test_util.run_v2_only
  def test_GRU_runtime(self):
    layer = rnn.GRU(self.rnn_state_size, return_runtime=True)

    inputs = keras.layers.Input(
        shape=[self.timestep, self.input_shape], dtype=dtypes.float32)

    outputs, runtime = layer(inputs)
    # Expand the runtime so that it is a 1D tensor instead of scalar.
    # TF model does not work with scalar model output, specially during
    # aggregation.
    runtime = keras.layers.Lambda(
        lambda x: array_ops.expand_dims(x, axis=-1))(runtime)
    model = keras.models.Model(inputs=inputs, outputs=[outputs, runtime])
    self._test_runtime_with_model(model)

  @test_util.run_v2_only
  def test_GRU_runtime_with_mask(self):
    # Masking will affect which backend is selected based on whether the mask
    # is strictly right padded.
    layer = rnn.GRU(self.rnn_state_size, return_runtime=True)

    inputs = keras.layers.Input(
        shape=[self.timestep, self.input_shape], dtype=dtypes.float32)
    masked_inputs = keras.layers.Masking()(inputs)

    outputs, runtime = layer(masked_inputs)
    # Expand the runtime so that it is a 1D tensor instead of scalar.
    # TF model does not work with scalar model output, specially during
    # aggregation.
    runtime = keras.layers.Lambda(
        lambda x: array_ops.expand_dims(x, axis=-1))(runtime)
    model = keras.models.Model(inputs=inputs, outputs=[outputs, runtime])

    (x_train, y_train), _ = testing_utils.get_test_data(
        train_samples=self.batch,
        test_samples=0,
        input_shape=(self.timestep, self.input_shape),
        num_classes=self.output_shape)
    y_train = keras.utils.to_categorical(y_train, self.output_shape)

    model.compile(
        optimizer='sgd',
        loss=['categorical_crossentropy', None],
        run_eagerly=testing_utils.should_run_eagerly(),
        experimental_run_tf_function=testing_utils.should_run_tf_function())

    model.fit(x_train, y_train)

    # Verify unpadded data.
    _, runtime_value = model.predict(x_train)
    if test.is_gpu_available():
      self.assertEqual(runtime_value[0], rnn._RUNTIME_GPU)
    else:
      self.assertEqual(runtime_value[0], rnn._RUNTIME_CPU)

    # Update x/y to be right padded by setting the last timestep to 0
    x_train[:, -1, :] = 0
    y_train[:, -1] = 0
    _, runtime_value = model.predict(x_train)
    if test.is_gpu_available():
      self.assertEqual(runtime_value[0], rnn._RUNTIME_GPU)
    else:
      self.assertEqual(runtime_value[0], rnn._RUNTIME_CPU)

    # Further update x/y to be mix padded (masks in the middle), and verify
    # only cpu kernel can be selected.
    x_train[:, -3, :] = 0
    y_train[:, -3] = 0
    _, runtime_value = model.predict(x_train)
    self.assertEqual(runtime_value[0], rnn._RUNTIME_CPU)

  @test_util.run_v2_only
  def test_GRU_runtime_with_cond(self):
    # This test is to demonstrate the graph rewrite of grappler plugin under
    # the condition that the function returns different number of internal
    # states.
    layer = rnn.GRU(self.rnn_state_size, return_runtime=True)

    inputs = keras.layers.Input(
        shape=[self.timestep, self.input_shape], dtype=dtypes.float32)

    zeros = array_ops.zeros([self.batch, self.output_shape])
    dummy_runtime = rnn._runtime(rnn._RUNTIME_UNKNOWN)
    a = constant_op.constant(0)
    b = constant_op.constant(1)
    # Will always run the GRU layer.
    outputs, runtime = control_flow_ops.cond(
        gen_math_ops.less(a, b),
        lambda: layer(inputs),
        lambda: (zeros, dummy_runtime))

    # Expand the runtime so that it is a 1D tensor instead of scalar.
    # TF model does not work with scalar model output, specially during
    # aggregation.
    runtime = keras.layers.Lambda(
        lambda x: array_ops.expand_dims(x, axis=-1))(runtime)
    model = keras.models.Model(inputs=inputs, outputs=[outputs, runtime])
    self._test_runtime_with_model(model)


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