# Copyright 2017 The Sonnet 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 recurrent cores in snt."""

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

import itertools

# Dependency imports
from absl.testing import parameterized
import numpy as np
from six.moves import xrange  # pylint: disable=redefined-builtin
import sonnet as snt
import tensorflow.compat.v1 as tf
from tensorflow.contrib import rnn as contrib_rnn
from tensorflow.contrib.eager.python import tfe as contrib_eager

from tensorflow.python.ops import variables  # pylint: disable=g-direct-tensorflow-import


# Some helpers used for generic tests which cover both LSTM and BatchNormLSTM:


def _get_lstm_variable_names(lstm):
  if isinstance(lstm, snt.BatchNormLSTM):
    var_names = lstm.get_possible_initializer_keys(
        lstm.use_peepholes,
        lstm.use_batch_norm_h,
        lstm.use_batch_norm_x,
        lstm.use_batch_norm_c)
    if lstm.use_batch_norm_h or lstm.use_batch_norm_x:
      var_names |= {"w_x", "w_h"}
    else:
      var_names |= {"w_xh"}
  else:
    var_names = lstm.get_possible_initializer_keys(lstm.use_peepholes)
    var_names |= {"w_xh"}
  var_names -= {"w_gates", "b_gates"}
  var_names |= {"b"}
  return var_names


def _construct_lstm(use_batch_norm_h=False, use_batch_norm_x=False,
                    use_batch_norm_c=False, max_unique_stats=1, **kwargs):
  if any([use_batch_norm_h, use_batch_norm_x, use_batch_norm_c]):
    cell = snt.BatchNormLSTM(
        use_batch_norm_h=use_batch_norm_h,
        use_batch_norm_x=use_batch_norm_x,
        use_batch_norm_c=use_batch_norm_c,
        max_unique_stats=max_unique_stats,
        **kwargs)
    return cell, cell.with_batch_norm_control(is_training=True)
  else:
    cell = snt.LSTM(**kwargs)
    return cell, cell


def _get_possible_initializer_keys(use_peepholes, use_batch_norm_h,
                                   use_batch_norm_x, use_batch_norm_c):
  if any([use_batch_norm_h, use_batch_norm_x, use_batch_norm_c]):
    return snt.BatchNormLSTM.get_possible_initializer_keys(
        use_peepholes, use_batch_norm_h, use_batch_norm_x, use_batch_norm_c)
  else:
    return snt.LSTM.get_possible_initializer_keys(use_peepholes)


@contrib_eager.run_all_tests_in_graph_and_eager_modes
class LSTMTest(tf.test.TestCase, parameterized.TestCase):

  def testShape(self):
    batch_size = 2
    hidden_size = 4
    inputs = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    prev_hidden = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    prev_cell = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    lstm = snt.LSTM(hidden_size)
    output, next_state = lstm(inputs, (prev_hidden, prev_cell))

    shape = np.ndarray((batch_size, hidden_size))

    self.assertShapeEqual(shape, next_state[0])
    self.assertShapeEqual(shape, next_state[1])
    self.assertShapeEqual(shape, output)

  def testVariables(self):
    batch_size = 5
    hidden_size = 20
    mod_name = "rnn"
    inputs = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    prev_cell = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    prev_hidden = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    lstm = snt.LSTM(hidden_size, name=mod_name)
    self.assertEqual(lstm.scope_name, mod_name)
    with self.assertRaisesRegexp(snt.Error, "not instantiated yet"):
      lstm.get_variables()
    lstm(inputs, (prev_hidden, prev_cell))

    lstm_variables = lstm.get_variables()
    self.assertLen(lstm_variables, 2, "LSTM should have 2 variables")
    param_map = {param.name.split("/")[-1].split(":")[0]:
                 param for param in lstm_variables}

    self.assertShapeEqual(
        np.ndarray(4 * hidden_size),
        tf.convert_to_tensor(param_map[snt.LSTM.B_GATES]))
    self.assertShapeEqual(
        np.ndarray((2 * hidden_size, 4 * hidden_size)),
        tf.convert_to_tensor(param_map[snt.LSTM.W_GATES]))

  @parameterized.named_parameters(
      [("lstm", None), ("lstm_with_recurrent_projection", 6)])
  def testComputation(self, projection_size):
    batch_size = 2
    hidden_size = 4
    hidden_state_size = projection_size or hidden_size

    # With random data, check the TF calculation matches the Numpy version.
    input_data = np.random.randn(batch_size, hidden_size).astype(np.float32)
    prev_hidden_data = np.random.randn(batch_size,
                                       hidden_state_size).astype(np.float32)
    prev_cell_data = np.random.randn(batch_size, hidden_size).astype(np.float32)

    inputs = tf.constant(input_data)
    prev_cell = tf.constant(prev_cell_data)
    prev_hidden = tf.constant(prev_hidden_data)

    lstm = snt.LSTM(hidden_size, projection_size=projection_size)
    _, next_state = lstm(inputs, (prev_hidden, prev_cell))
    next_hidden, next_cell = next_state
    lstm_variables = lstm.get_variables()
    param_map = {param.name.split("/")[-1].split(":")[0]:
                 param for param in lstm_variables}

    self.evaluate(tf.global_variables_initializer())
    fetches = [(next_hidden, next_cell), param_map[snt.LSTM.W_GATES],
               param_map[snt.LSTM.B_GATES]]
    if projection_size is not None:
      fetches.append(param_map[snt.LSTM.W_H_PROJECTION])
    output = self.evaluate(fetches)

    next_state_ex, gate_weights_ex, gate_biases_ex = output[:3]
    in_and_hid = np.concatenate((input_data, prev_hidden_data), axis=1)
    real_gate = np.dot(in_and_hid, gate_weights_ex) + gate_biases_ex
    # i = input_gate, j = next_input, f = forget_gate, o = output_gate
    i, j, f, o = np.hsplit(real_gate, 4)
    real_cell = (prev_cell_data / (1 + np.exp(-(f + lstm._forget_bias))) +
                 1 / (1 + np.exp(-i)) * np.tanh(j))
    real_hidden = np.tanh(real_cell) * 1 / (1 + np.exp(-o))
    if projection_size is not None:
      real_hidden = np.matmul(real_hidden, output[-1])

    self.assertAllClose(real_hidden, next_state_ex[0])
    self.assertAllClose(real_cell, next_state_ex[1])

  def testPeephole(self):
    batch_size = 5
    hidden_size = 20

    # With random data, check the TF calculation matches the Numpy version.
    input_data = np.random.randn(batch_size, hidden_size).astype(np.float32)
    prev_hidden_data = np.random.randn(batch_size,
                                       hidden_size).astype(np.float32)
    prev_cell_data = np.random.randn(batch_size, hidden_size).astype(np.float32)

    # Initialize the rnn and verify the number of parameter sets.
    inputs = tf.constant(input_data)
    prev_cell = tf.constant(prev_cell_data)
    prev_hidden = tf.constant(prev_hidden_data)

    lstm = snt.LSTM(hidden_size, use_peepholes=True)
    _, next_state = lstm(inputs, (prev_hidden, prev_cell))
    next_hidden, next_cell = next_state
    lstm_variables = lstm.get_variables()
    self.assertLen(lstm_variables, 5, "LSTM should have 5 variables")

    # Unpack parameters into dict and check their sizes.
    param_map = {param.name.split("/")[-1].split(":")[0]:
                 param for param in lstm_variables}
    self.assertShapeEqual(
        np.ndarray(4 * hidden_size),
        tf.convert_to_tensor(param_map[snt.LSTM.B_GATES]))
    self.assertShapeEqual(
        np.ndarray((2 * hidden_size, 4 * hidden_size)),
        tf.convert_to_tensor(param_map[snt.LSTM.W_GATES]))
    self.assertShapeEqual(
        np.ndarray(hidden_size),
        tf.convert_to_tensor(param_map[snt.LSTM.W_F_DIAG]))
    self.assertShapeEqual(
        np.ndarray(hidden_size),
        tf.convert_to_tensor(param_map[snt.LSTM.W_I_DIAG]))
    self.assertShapeEqual(
        np.ndarray(hidden_size),
        tf.convert_to_tensor(param_map[snt.LSTM.W_O_DIAG]))

    self.evaluate(tf.global_variables_initializer())
    fetches = [(next_hidden, next_cell), param_map[snt.LSTM.W_GATES],
               param_map[snt.LSTM.B_GATES], param_map[snt.LSTM.W_F_DIAG],
               param_map[snt.LSTM.W_I_DIAG], param_map[snt.LSTM.W_O_DIAG]]
    output = self.evaluate(fetches)

    next_state_ex, w_ex, b_ex, wfd_ex, wid_ex, wod_ex = output
    in_and_hid = np.concatenate((input_data, prev_hidden_data), axis=1)
    real_gate = np.dot(in_and_hid, w_ex) + b_ex
    # i = input_gate, j = next_input, f = forget_gate, o = output_gate
    i, j, f, o = np.hsplit(real_gate, 4)
    real_cell = (prev_cell_data /
                 (1 + np.exp(-(f + lstm._forget_bias +
                               wfd_ex * prev_cell_data))) +
                 1 / (1 + np.exp(-(i + wid_ex * prev_cell_data))) * np.tanh(j))
    real_hidden = (np.tanh(real_cell + wod_ex * real_cell) *
                   1 / (1 + np.exp(-o)))

    self.assertAllClose(real_hidden, next_state_ex[0])
    self.assertAllClose(real_cell, next_state_ex[1])

  @parameterized.parameters(
      *itertools.product(
          (True, False), (True, False), (True, False), (True, False))
  )
  def testInitializers(self, use_peepholes, use_batch_norm_h, use_batch_norm_x,
                       use_batch_norm_c):
    batch_size = 2
    hidden_size = 4

    keys = _get_possible_initializer_keys(
        use_peepholes, use_batch_norm_h, use_batch_norm_x, use_batch_norm_c)
    initializers = {
        key: tf.constant_initializer(1.5) for key in keys
    }

    # Test we can successfully create the LSTM with initializers.
    lstm, wrapped_lstm = _construct_lstm(hidden_size=hidden_size,
                                         use_peepholes=use_peepholes,
                                         use_batch_norm_h=use_batch_norm_h,
                                         use_batch_norm_x=use_batch_norm_x,
                                         use_batch_norm_c=use_batch_norm_c,
                                         initializers=initializers)

    # Test we can build the LSTM.
    inputs = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    prev_cell = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    prev_hidden = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    wrapped_lstm(inputs, (prev_hidden, prev_cell))
    init = tf.global_variables_initializer()

    # Test that the initializers have been correctly applied.
    lstm_variable_names = _get_lstm_variable_names(lstm)
    lstm_variables = [getattr(lstm, "_" + name) for name in lstm_variable_names]

    self.evaluate(init)
    lstm_variables_v = self.evaluate(lstm_variables)
    for lstm_variable_v in lstm_variables_v:
      self.assertAllClose(lstm_variable_v, 1.5 * np.ones(lstm_variable_v.shape))

  def testPeepholeInitializersCheck(self):
    hidden_size = 4

    # Test that passing in a peephole initializer when we don't request peephole
    # connections raises an error.
    for key in [snt.LSTM.W_F_DIAG, snt.LSTM.W_I_DIAG, snt.LSTM.W_O_DIAG]:
      with self.assertRaisesRegexp(KeyError, "Invalid initializer"):
        snt.LSTM(hidden_size, use_peepholes=False,
                 initializers={key: tf.constant_initializer(0)})

  @parameterized.parameters(
      (1e-6, 14),
      (0.5, None),
      (1 - 1e-6, 0)
  )
  def testRecurrentDropout(self, keep_prob, expected_zeros):
    """Performs various recurrent dropout checks.

    - The training and testing versions have the same output when the keep prob
      is very close to 1.
    - The return is deterministic for keep probs 0 or close to 1.
    - The final hidden state has 0s at the same position as the mask.

    Args:
      keep_prob: the recurrent dropout keep probability.
      expected_zeros: the number of expected zeros in the dropout mask.
    """

    batch_size = 2
    input_size = 3
    hidden_size = 7
    seq_len = 5

    train_cell, test_cell = snt.lstm_with_recurrent_dropout(
        hidden_size, keep_prob=keep_prob)
    input_data = np.stack(
        [np.random.rand(seq_len, input_size)] * batch_size).astype(np.float32)
    inputs = tf.constant(input_data)
    train_output, ((train_hidden, _), [train_mask]) = tf.nn.dynamic_rnn(
        train_cell,
        inputs,
        initial_state=train_cell.initial_state(batch_size, tf.float32),
        dtype=tf.float32)
    valid_output, _ = tf.nn.dynamic_rnn(
        test_cell,
        inputs,
        initial_state=test_cell.initial_state(
            batch_size, tf.float32),
        dtype=tf.float32)

    self.evaluate(tf.global_variables_initializer())
    # Use the same input data for each row.
    input_data = np.stack([np.random.rand(seq_len, input_size)] * batch_size)

    train_out, valid_out, hidden, mask = self.evaluate(
        [train_output, valid_output, train_hidden, train_mask])
    almost_one = abs(1 - keep_prob) < 1e-5
    if almost_one:
      self.assertAllClose(train_out, valid_out)
    else:
      self.assertGreater(np.max(train_out - valid_out), 0.001)

    self.assertAllClose(valid_out[0], valid_out[1])
    deterministic = almost_one or abs(keep_prob < 1e-5)
    if deterministic:
      self.assertAllClose(train_out[0], train_out[1])
      self.assertEqual(expected_zeros, np.sum(hidden == 0))
      self.assertEqual(expected_zeros, np.sum(mask == 0))
    else:
      self.assertGreater(np.max(train_out[0] - train_out[1]), 0.001)

    self.assertAllEqual(mask == 0, hidden == 0)

  @parameterized.parameters(
      (1 - 1e-6, 0, 0),
      (0.5, None, None),
      (1e-6, 14, 14)
  )
  def testZoneout(self, keep_prob, expected_frozen_h, expected_frozen_c):
    """Performs various zoneout checks.

    The training and testing versions have the same output when keep_prob
    is close to 0 or close to 1. The returned output for the training version
    is also deterministic in this case.

    Args:
      keep_prob: the probability to use the updated version of the state.
      expected_frozen_h: the number of hidden state values that are left
        unchanged after applying one step of LSTM with zoneout.
      expected_frozen_c: the number of cell state values that are left
        unchanged after applying one step of LSTM with zoneout.
    """

    batch_size = 2
    input_size = 3
    hidden_size = 7
    seq_len = 5

    train_cell, test_cell = snt.lstm_with_zoneout(
        hidden_size, keep_prob_c=keep_prob, keep_prob_h=keep_prob)
    # Use the same input data for each row.
    input_data = np.stack(
        [np.random.rand(seq_len, input_size)] * batch_size).astype(np.float32)
    inputs = tf.constant(input_data)
    train_output, (train_h, train_c) = tf.nn.dynamic_rnn(
        train_cell,
        inputs,
        dtype=tf.float32)
    _, (next_train_h, next_train_c) = train_cell(
        inputs[:, 0], snt.LSTMState(train_h, train_c))
    valid_output, _ = tf.nn.dynamic_rnn(
        test_cell,
        inputs,
        dtype=tf.float32)

    self.evaluate(tf.global_variables_initializer())

    outputs = self.evaluate({
        "train_out": train_output,
        "valid_out": valid_output,
        "train_h": train_h,
        "train_c": train_c,
        "next_train_h": next_train_h,
        "next_train_c": next_train_c,
    })

    self.assertAllClose(outputs["valid_out"][0], outputs["valid_out"][1])
    deterministic = abs(1 - keep_prob) < 1e-5 or abs(keep_prob < 1e-5)
    if deterministic:
      self.assertAllClose(outputs["train_out"], outputs["valid_out"])
      self.assertAllClose(outputs["train_out"][0], outputs["train_out"][1])
      self.assertEqual(expected_frozen_h,
                       np.sum(outputs["train_h"] == outputs["next_train_h"]))
      self.assertEqual(expected_frozen_c,
                       np.sum(outputs["train_c"] == outputs["next_train_c"]))
    else:
      # Ensure that the training and validation outputs are different.
      self.assertGreater(
          np.max(np.abs(outputs["train_out"] - outputs["valid_out"])), 0.05)
      # Ensure that the training output is not deterministic.
      self.assertGreater(
          np.max(np.abs(outputs["train_out"][0] - outputs["train_out"][1])),
          0.04)

  @parameterized.parameters(
      (True, False, False),
      (False, True, False),
      (False, False, True)
  )
  def testBatchNormBuildFlag(self, use_batch_norm_h, use_batch_norm_x,
                             use_batch_norm_c):
    """Check if an error is raised if we don't specify the is_training flag."""
    batch_size = 2
    hidden_size = 4

    inputs = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    prev_cell = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    prev_hidden = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])

    err = "is_training flag must be explicitly specified"
    with self.assertRaisesRegexp(ValueError, err):
      lstm = snt.BatchNormLSTM(
          hidden_size,
          use_batch_norm_h=use_batch_norm_h,
          use_batch_norm_x=use_batch_norm_x,
          use_batch_norm_c=use_batch_norm_c)
      lstm(inputs, (prev_cell, prev_hidden))

  def testBatchNormInitializersCheck(self):
    hidden_size = 4

    # Test that passing in a batchnorm initializer when we don't request
    # that form of batchnorm raises an error.
    for key, options in [
        (snt.BatchNormLSTM.GAMMA_H, {"use_batch_norm_h": False,
                                     "use_batch_norm_x": True}),
        (snt.BatchNormLSTM.GAMMA_X, {"use_batch_norm_x": False,
                                     "use_batch_norm_h": True}),
        (snt.BatchNormLSTM.GAMMA_C, {"use_batch_norm_c": False,
                                     "use_batch_norm_h": True}),
        (snt.BatchNormLSTM.BETA_C, {"use_batch_norm_c": False,
                                    "use_batch_norm_h": True})]:
      with self.assertRaisesRegexp(KeyError, "Invalid initializer"):
        snt.BatchNormLSTM(
            hidden_size, initializers={key: tf.constant_initializer(0)},
            **options)

  @parameterized.parameters(
      *itertools.product(
          (True, False), (True, False), (True, False), (True, False))
  )
  def testPartitioners(self, use_peepholes, use_batch_norm_h, use_batch_norm_x,
                       use_batch_norm_c):
    if tf.executing_eagerly():
      self.skipTest("Partitioned variables arenot supported in eager mode.")
    batch_size = 2
    hidden_size = 4

    keys = _get_possible_initializer_keys(
        use_peepholes, use_batch_norm_h, use_batch_norm_x, use_batch_norm_c)
    partitioners = {
        key: tf.variable_axis_size_partitioner(10) for key in keys
    }

    # Test we can successfully create the LSTM with partitioners.
    lstm, wrapped_lstm = _construct_lstm(hidden_size=hidden_size,
                                         use_peepholes=use_peepholes,
                                         use_batch_norm_h=use_batch_norm_h,
                                         use_batch_norm_x=use_batch_norm_x,
                                         use_batch_norm_c=use_batch_norm_c,
                                         partitioners=partitioners)

    # Test we can build the LSTM
    inputs = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    prev_cell = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    prev_hidden = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    wrapped_lstm(inputs, (prev_hidden, prev_cell))

    # Test that the variables are partitioned.
    var_names = _get_lstm_variable_names(lstm)
    for var_name in var_names:
      self.assertEqual(type(getattr(lstm, "_" + var_name)),
                       variables.PartitionedVariable)

  @parameterized.parameters(
      *itertools.product(
          (True, False), (True, False), (True, False), (True, False))
  )
  def testRegularizers(self, use_peepholes, use_batch_norm_h, use_batch_norm_x,
                       use_batch_norm_c):
    batch_size = 2
    hidden_size = 4

    keys = _get_possible_initializer_keys(
        use_peepholes, use_batch_norm_h, use_batch_norm_x, use_batch_norm_c)
    regularizers = {
        key: tf.nn.l2_loss for key in keys
    }

    # Test we can successfully create the LSTM with regularizers.
    _, wrapped_lstm = _construct_lstm(hidden_size=hidden_size,
                                      use_peepholes=use_peepholes,
                                      use_batch_norm_h=use_batch_norm_h,
                                      use_batch_norm_x=use_batch_norm_x,
                                      use_batch_norm_c=use_batch_norm_c,
                                      regularizers=regularizers)

    # Test we can build the LSTM
    inputs = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    prev_cell = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    prev_hidden = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    wrapped_lstm(inputs, (prev_hidden, prev_cell))

    # Test that we have regularization losses.
    num_reg_losses = len(tf.get_collection(
        tf.GraphKeys.REGULARIZATION_LOSSES))
    if use_batch_norm_h or use_batch_norm_x:
      self.assertEqual(num_reg_losses, len(keys) + 1)
    else:
      self.assertLen(keys, num_reg_losses)

  # Pick some hopefully representative combination of parameter values
  # (want to test with seq_len < max_unique_stats and seq_len >
  # max_unique_stats, and some other combinations for good measure).
  @parameterized.parameters(
      (False, 1, 1, 2),
      (True, 3, 1, 2),
      (False, 1, 2, 1),
      (True, 3, 2, 1),
      (False, 3, 3, 5))
  def testSameInStaticAndDynamicWithBatchNorm(self, use_peepholes, batch_size,
                                              max_unique_stats, seq_len):
    # Tests that when the cell is used in either a normal tensorflow rnn, or in
    # tensorflow's dynamic_rnn, that the output is the same. This is to test
    # test that the cores aren't doing anything funny they shouldn't be (like
    # relying on the number of times they've been invoked).

    hidden_size = 3
    input_size = 3

    inputs = tf.ones(
        dtype=tf.float32,
        shape=[batch_size, seq_len, input_size],
        name="inputs")
    static_inputs = tf.unstack(inputs, axis=1)

    test_local_stats = False

    cell = snt.BatchNormLSTM(
        hidden_size=hidden_size,
        max_unique_stats=max_unique_stats,
        use_peepholes=use_peepholes,
        use_batch_norm_h=True,
        use_batch_norm_x=True,
        use_batch_norm_c=True)

    def connect(training):
      static_output_unpacked, _ = contrib_rnn.static_rnn(
          cell.with_batch_norm_control(
              is_training=training, test_local_stats=test_local_stats),
          static_inputs,
          initial_state=cell.initial_state(batch_size, tf.float32))

      static_output = tf.stack(static_output_unpacked, axis=1)

      dynamic_output, _ = tf.nn.dynamic_rnn(
          cell.with_batch_norm_control(is_training=training,
                                       test_local_stats=test_local_stats),
          inputs,
          initial_state=cell.initial_state(batch_size, tf.float32),
          dtype=tf.float32)

      return static_output, dynamic_output

    ops = connect(training=True)
    self.evaluate(tf.global_variables_initializer())
    static_out, dynamic_out = self.evaluate(ops)
    self.assertAllClose(static_out, dynamic_out)

    # Do a pass to train the exponential moving statistics.
    for _ in range(5):
      if tf.executing_eagerly():
        ops = connect(training=True)
      static_out, dynamic_out = self.evaluate(ops)
      self.assertAllClose(static_out, dynamic_out)

    # And check that same when using test statistics.
    ops = connect(training=False)
    static_out, dynamic_out = self.evaluate(ops)
    self.assertAllClose(static_out, dynamic_out)

  def testSameInStaticAndDynamic(self):
    batch_size = 3
    seq_len = 2
    hidden_size = 3
    input_size = 3

    inputs = tf.ones(
        dtype=tf.float32,
        shape=[batch_size, seq_len, input_size],
        name="inputs")
    static_inputs = tf.unstack(inputs, axis=1)

    cell = snt.LSTM(hidden_size=hidden_size)

    static_output_unpacked, _ = contrib_rnn.static_rnn(
        cell,
        static_inputs,
        initial_state=cell.initial_state(batch_size, tf.float32))

    dynamic_output, _ = tf.nn.dynamic_rnn(
        cell, inputs,
        initial_state=cell.initial_state(batch_size, tf.float32),
        dtype=tf.float32)

    static_output = tf.stack(static_output_unpacked, axis=1)

    self.evaluate(tf.global_variables_initializer())

    # Check that static and dynamic give the same output
    static_out, dynamic_out = self.evaluate([static_output, dynamic_output])
    self.assertAllClose(static_out, dynamic_out)

  def testLayerNormVariables(self):
    core = snt.LSTM(hidden_size=3, use_layer_norm=True)

    batch_size = 3
    inputs = tf.ones(dtype=tf.float32, shape=[batch_size, 3, 3])
    tf.nn.dynamic_rnn(core,
                      inputs,
                      initial_state=core.initial_state(batch_size, tf.float32))

    self.assertTrue(core.use_layer_norm)

    expected = 4  # gate bias and one weight, plus LayerNorm's gamma, beta.
    self.assertLen(core.get_variables(), expected)

  def testHiddenClipping(self):
    core = snt.LSTM(hidden_size=5, hidden_clip_value=1.0)
    obs = tf.constant(np.random.rand(3, 10), dtype=tf.float32)

    unclipped = np.random.rand(3, 5) - 0.5
    unclipped *= 2.0 / unclipped.max()
    unclipped = unclipped.astype(np.float32)
    clipped = unclipped.clip(-1., 1.)

    hidden = tf.constant(unclipped)
    cell = tf.constant(unclipped)
    output = core(obs, [hidden, cell])
    self.evaluate(tf.global_variables_initializer())
    output1, (hidden1, cell1) = self.evaluate(output)

    hidden = tf.constant(clipped)
    output = core(obs, [hidden, cell])

    output2, (hidden2, cell2) = self.evaluate(output)
    self.assertAllClose(output1, output2)
    self.assertAllClose(hidden1, hidden2)
    self.assertAllClose(cell1, cell2)

  def testCellClipping(self):
    core = snt.LSTM(hidden_size=5, cell_clip_value=1.0)
    obs = tf.constant(np.random.rand(3, 10), dtype=tf.float32)

    unclipped = np.random.rand(3, 5) - 0.5
    unclipped *= 2.0 / unclipped.max()
    unclipped = unclipped.astype(np.float32)
    clipped = unclipped.clip(-1., 1.)

    hidden = tf.constant(unclipped)
    cell = tf.constant(unclipped)
    output = core(obs, [hidden, cell])
    self.evaluate(tf.global_variables_initializer())
    output1, (hidden1, cell1) = self.evaluate(output)

    cell = tf.constant(clipped)
    output = core(obs, [hidden, cell])
    output2, (hidden2, cell2) = self.evaluate(output)
    self.assertAllClose(output1, output2)
    self.assertAllClose(hidden1, hidden2)
    self.assertAllClose(cell1, cell2)

  @parameterized.parameters(
      (False, False, False, False),
      (False, True, False, False),
      (True, False, True, False),
      (False, True, True, False),
      (False, False, False, True),
      (True, True, False, True),
      (False, False, True, True),
      (False, True, True, True))
  def testBatchNormVariables(self,
                             use_peepholes,
                             use_batch_norm_h,
                             use_batch_norm_x,
                             use_batch_norm_c):
    cell, wrapped_cell = _construct_lstm(hidden_size=3,
                                         use_peepholes=use_peepholes,
                                         use_batch_norm_h=use_batch_norm_h,
                                         use_batch_norm_x=use_batch_norm_x,
                                         use_batch_norm_c=use_batch_norm_c)

    # Need to connect the cell before it has variables
    batch_size = 3
    inputs = tf.ones(dtype=tf.float32, shape=[batch_size, 3, 3])
    tf.nn.dynamic_rnn(wrapped_cell, inputs,
                      initial_state=cell.initial_state(batch_size, tf.float32))

    self.assertEqual(use_peepholes, cell.use_peepholes)
    if use_batch_norm_h or use_batch_norm_x or use_batch_norm_c:
      self.assertEqual(use_batch_norm_h, cell.use_batch_norm_h)
      self.assertEqual(use_batch_norm_x, cell.use_batch_norm_x)
      self.assertEqual(use_batch_norm_c, cell.use_batch_norm_c)

    if use_batch_norm_h or use_batch_norm_x:
      expected = 3  # gate bias and two weights
    else:
      expected = 2  # gate bias and weight
    if use_peepholes:
      expected += 3
    if use_batch_norm_h:
      expected += 1  # gamma_h
    if use_batch_norm_x:
      expected += 1  # gamma_x
    if use_batch_norm_c:
      expected += 2  # gamma_c, beta_c

    self.assertLen(cell.get_variables(), expected)

  def testCheckMaxUniqueStats(self):
    self.assertRaisesRegexp(ValueError,
                            ".*must be >= 1",
                            snt.BatchNormLSTM,
                            hidden_size=1,
                            max_unique_stats=0)

  @parameterized.parameters(
      (False, 1),
      (False, 2),
      (True, 1),
      (True, 2))
  def testTraining(self, trainable_initial_state, max_unique_stats):
    """Test that everything trains OK, with or without trainable init. state."""
    hidden_size = 3
    batch_size = 3
    time_steps = 3
    cell = snt.BatchNormLSTM(hidden_size=hidden_size,
                             max_unique_stats=max_unique_stats)
    inputs = tf.constant(np.random.rand(batch_size, time_steps, 3),
                         dtype=tf.float32)
    initial_state = cell.initial_state(
        batch_size, tf.float32, trainable_initial_state)
    def loss_fn():
      output, _ = tf.nn.dynamic_rnn(
          cell.with_batch_norm_control(is_training=True),
          inputs,
          initial_state=initial_state,
          dtype=tf.float32)

      return tf.reduce_mean(tf.square(
          output - np.random.rand(batch_size, time_steps, hidden_size)))

    train_op = tf.train.GradientDescentOptimizer(1).minimize(
        loss_fn if tf.executing_eagerly() else loss_fn())

    init = tf.global_variables_initializer()
    self.evaluate(init)
    self.evaluate(train_op)

  # Regression test.
  def testSideBySide(self):
    hidden_size = 3
    batch_size = 4
    lstm1 = snt.LSTM(hidden_size=hidden_size)
    lstm2 = snt.LSTM(hidden_size=hidden_size)
    lstm1.initial_state(batch_size, trainable=True)
    # Previously either of the two lines below would cause a crash due to
    # Variable name collision.
    lstm1.initial_state(batch_size, trainable=True)
    lstm2.initial_state(batch_size, trainable=True)

  def testInitialStateNames(self):
    if tf.executing_eagerly():
      self.skipTest("Tensor name is not supported in eager mode.")
    lstm = snt.LSTM(hidden_size=3, name="foo")
    unnamed_init_state = lstm.initial_state(4, trainable=True)
    named_init_state = lstm.initial_state(4, trainable=True, name="bar")
    self.assertEqual(unnamed_init_state[0].name,
                     "foo_initial_state/state_hidden_tiled:0")
    self.assertEqual(unnamed_init_state[1].name,
                     "foo_initial_state/state_cell_tiled:0")
    self.assertEqual(named_init_state[0].name, "bar/state_hidden_tiled:0")
    self.assertEqual(named_init_state[1].name, "bar/state_cell_tiled:0")


@contrib_eager.run_all_tests_in_graph_and_eager_modes
class ConvLSTMTest(tf.test.TestCase, parameterized.TestCase):

  @parameterized.parameters(
      (snt.Conv1DLSTM, 1, False),
      (snt.Conv1DLSTM, 1, True),
      (snt.Conv2DLSTM, 2, False),
      (snt.Conv2DLSTM, 2, True),
  )
  def testShape(self, lstm_class, dim, use_bias):
    batch_size = 2
    input_shape = (8,) * dim
    input_channels = 3
    output_channels = 5

    input_shape = (batch_size,) + input_shape + (input_channels,)
    output_shape = input_shape[:-1] + (output_channels,)

    inputs = tf.ones(dtype=tf.float32, shape=input_shape)
    prev_hidden = tf.ones(dtype=tf.float32, shape=output_shape)
    prev_cell = tf.ones(dtype=tf.float32, shape=output_shape)
    lstm = lstm_class(
        input_shape=input_shape[1:],
        output_channels=output_channels,
        kernel_shape=1,
        use_bias=use_bias)
    output, next_state = lstm(inputs, (prev_hidden, prev_cell))

    expected_shape = np.ndarray(output_shape)

    self.assertShapeEqual(expected_shape, next_state[0])
    self.assertShapeEqual(expected_shape, next_state[1])
    self.assertShapeEqual(expected_shape, output)

  @parameterized.parameters(
      (snt.Conv1DLSTM, 1, False, False),
      (snt.Conv1DLSTM, 1, True, False),
      (snt.Conv2DLSTM, 2, False, False),
      (snt.Conv2DLSTM, 2, True, False),
      (snt.Conv1DLSTM, 1, False, True),
      (snt.Conv1DLSTM, 1, True, True),
      (snt.Conv2DLSTM, 2, False, True),
      (snt.Conv2DLSTM, 2, True, True),
  )
  def testInitializers(self, lstm_class, dim, use_bias, legacy_bias_behaviour):
    keys = snt.Conv2DLSTM.get_possible_initializer_keys(use_bias)
    initializers = {
        key: tf.constant_initializer(i) for i, key in enumerate(keys)
    }

    batch_size = 2
    input_shape = (8,) * dim
    input_channels = 3
    output_channels = 5

    input_shape = (batch_size,) + input_shape + (input_channels,)
    output_shape = input_shape[:-1] + (output_channels,)

    inputs = tf.ones(dtype=tf.float32, shape=input_shape)
    prev_hidden = tf.ones(dtype=tf.float32, shape=output_shape)
    prev_cell = tf.ones(dtype=tf.float32, shape=output_shape)

    # Test we can successfully create the LSTM with partitioners.
    lstm = lstm_class(
        input_shape=input_shape[1:],
        output_channels=output_channels,
        kernel_shape=1,
        use_bias=use_bias,
        legacy_bias_behaviour=legacy_bias_behaviour,
        initializers=initializers)
    lstm(inputs, (prev_hidden, prev_cell))

    init = tf.global_variables_initializer()

    # Test that the initializers have been applied correctly.
    self.evaluate(init)
    for conv_key, convolution in lstm.convolutions.items():
      for i, key in enumerate(keys):
        if not legacy_bias_behaviour and conv_key == "hidden" and key == "b":
          self.assertFalse(hasattr(convolution, key))
          continue
        variable = getattr(convolution, key)
        self.assertAllClose(self.evaluate(variable),
                            np.full(variable.get_shape(),
                                    i, dtype=np.float32))

  @parameterized.parameters(
      (snt.Conv1DLSTM, 1, False, False),
      (snt.Conv1DLSTM, 1, True, False),
      (snt.Conv2DLSTM, 2, False, False),
      (snt.Conv2DLSTM, 2, True, False),
      (snt.Conv1DLSTM, 1, False, True),
      (snt.Conv1DLSTM, 1, True, True),
      (snt.Conv2DLSTM, 2, False, True),
      (snt.Conv2DLSTM, 2, True, True),
  )
  def testPartitioners(self, lstm_class, dim, use_bias, legacy_bias_behaviour):
    if tf.executing_eagerly():
      self.skipTest("Partitioned variables are not supported in eager.")
    keys = snt.Conv2DLSTM.get_possible_initializer_keys(use_bias)
    partitioners = {
        key: tf.variable_axis_size_partitioner(10) for key in keys
    }

    batch_size = 2
    input_shape = (8,) * dim
    input_channels = 3
    output_channels = 5

    input_shape = (batch_size,) + input_shape + (input_channels,)
    output_shape = input_shape[:-1] + (output_channels,)

    inputs = tf.ones(dtype=tf.float32, shape=input_shape)
    prev_hidden = tf.ones(dtype=tf.float32, shape=output_shape)
    prev_cell = tf.ones(dtype=tf.float32, shape=output_shape)

    # Test we can successfully create the LSTM with partitioners.
    lstm = lstm_class(
        input_shape=input_shape[1:],
        output_channels=output_channels,
        kernel_shape=1,
        use_bias=use_bias,
        legacy_bias_behaviour=legacy_bias_behaviour,
        partitioners=partitioners)
    lstm(inputs, (prev_hidden, prev_cell))

    # Test that the variables are partitioned.
    for conv_key, convolution in lstm.convolutions.items():
      for key in keys:
        if not legacy_bias_behaviour and conv_key == "hidden" and key == "b":
          self.assertFalse(hasattr(convolution, key))
          continue
        self.assertEqual(type(getattr(convolution, key)),
                         variables.PartitionedVariable)

  @parameterized.parameters(
      (snt.Conv1DLSTM, 1, False, False),
      (snt.Conv1DLSTM, 1, True, False),
      (snt.Conv2DLSTM, 2, False, False),
      (snt.Conv2DLSTM, 2, True, False),
      (snt.Conv1DLSTM, 1, False, True),
      (snt.Conv1DLSTM, 1, True, True),
      (snt.Conv2DLSTM, 2, False, True),
      (snt.Conv2DLSTM, 2, True, True),
  )
  def testRegularizers(self, lstm_class, dim, use_bias, legacy_bias_behaviour):
    keys = snt.Conv2DLSTM.get_possible_initializer_keys(use_bias)

    batch_size = 2
    input_shape = (8,) * dim
    input_channels = 3
    output_channels = 5

    input_shape = (batch_size,) + input_shape + (input_channels,)
    output_shape = input_shape[:-1] + (output_channels,)

    inputs = tf.ones(dtype=tf.float32, shape=input_shape)
    prev_hidden = tf.ones(dtype=tf.float32, shape=output_shape)
    prev_cell = tf.ones(dtype=tf.float32, shape=output_shape)

    # Test we can successfully create the LSTM with partitioners.
    lstm = lstm_class(
        input_shape=input_shape[1:],
        output_channels=output_channels,
        kernel_shape=1,
        use_bias=use_bias,
        legacy_bias_behaviour=legacy_bias_behaviour,
        regularizers={key: tf.nn.l2_loss for key in keys})
    lstm(inputs, (prev_hidden, prev_cell))

    # Test that we have regularization losses.
    num_reg_losses = len(tf.get_collection(
        tf.GraphKeys.REGULARIZATION_LOSSES))
    num_reg_losses_expected = len(lstm.convolutions) * len(keys)
    if use_bias and not legacy_bias_behaviour:
      # Bias is not applied to hidden
      num_reg_losses_expected -= 1
    self.assertEqual(num_reg_losses, num_reg_losses_expected)

  @parameterized.parameters(
      (snt.Conv1DLSTM, 1, False),
      (snt.Conv1DLSTM, 1, True),
      (snt.Conv2DLSTM, 2, False),
      (snt.Conv2DLSTM, 2, True),
  )
  def testTraining(self, lstm_class, dim, trainable_initial_state):
    """Test that training works, with or without trainable initial state."""
    time_steps = 1
    batch_size = 2
    input_shape = (8,) * dim
    input_channels = 3
    output_channels = 5

    input_shape = (batch_size,) + input_shape + (input_channels,)

    lstm = lstm_class(
        input_shape=input_shape[1:],
        output_channels=output_channels,
        kernel_shape=1)
    inputs = tf.random_normal((time_steps,) + input_shape, dtype=tf.float32)
    initial_state = lstm.initial_state(
        batch_size, tf.float32, trainable_initial_state)

    def loss_fn():
      output, _ = tf.nn.dynamic_rnn(lstm,
                                    inputs,
                                    time_major=True,
                                    initial_state=initial_state,
                                    dtype=tf.float32)

      return tf.reduce_mean(tf.square(output))

    train_op = tf.train.GradientDescentOptimizer(1).minimize(
        loss_fn if tf.executing_eagerly() else loss_fn())
    init = tf.global_variables_initializer()
    self.evaluate(init)
    self.evaluate(train_op)

  @parameterized.parameters(
      (snt.Conv1DLSTM, 1, False, 1, 1),
      (snt.Conv1DLSTM, 1, False, 1, 5),
      (snt.Conv1DLSTM, 1, False, 6, 1),
      (snt.Conv1DLSTM, 1, False, 6, 5),
      (snt.Conv1DLSTM, 1, True, 1, 1),
      (snt.Conv1DLSTM, 1, True, 1, 5),
      (snt.Conv1DLSTM, 1, True, 6, 1),
      (snt.Conv1DLSTM, 1, True, 6, 5),
      (snt.Conv2DLSTM, 2, False, 1, 1),
      (snt.Conv2DLSTM, 2, False, 1, 5),
      (snt.Conv2DLSTM, 2, False, 6, 1),
      (snt.Conv2DLSTM, 2, False, 6, 5),
      (snt.Conv2DLSTM, 2, True, 1, 1),
      (snt.Conv2DLSTM, 2, True, 1, 5),
      (snt.Conv2DLSTM, 2, True, 6, 1),
      (snt.Conv2DLSTM, 2, True, 6, 5),
  )
  def testDilatedConv(self, lstm_class, dim, trainable_initial_state, rate,
                      kernel_shape):
    """Test that training works, with or without dilated convolutions."""
    time_steps = 1
    batch_size = 2
    input_shape = (8,) * dim
    input_channels = 3
    output_channels = 5

    input_shape = (batch_size,) + input_shape + (input_channels,)

    lstm = lstm_class(
        input_shape=input_shape[1:],
        output_channels=output_channels,
        kernel_shape=kernel_shape,
        rate=rate)
    inputs = tf.random_normal((time_steps,) + input_shape, dtype=tf.float32)
    initial_state = lstm.initial_state(
        batch_size, tf.float32, trainable_initial_state)

    def loss_fn():
      output, _ = tf.nn.dynamic_rnn(lstm,
                                    inputs,
                                    time_major=True,
                                    initial_state=initial_state,
                                    dtype=tf.float32)

      return tf.reduce_mean(tf.square(output))

    train_op = tf.train.GradientDescentOptimizer(1).minimize(
        loss_fn if tf.executing_eagerly() else loss_fn())
    init = tf.global_variables_initializer()
    self.evaluate(init)
    self.evaluate(train_op)

  @parameterized.parameters(
      (snt.Conv1DLSTM, 1, False, False, 2),
      (snt.Conv1DLSTM, 1, False, True, 4),
      (snt.Conv1DLSTM, 1, True, False, 4),
      (snt.Conv1DLSTM, 1, True, True, 6),
      (snt.Conv2DLSTM, 2, False, False, 2),
      (snt.Conv2DLSTM, 2, False, True, 4),
      (snt.Conv2DLSTM, 2, True, False, 4),
      (snt.Conv2DLSTM, 2, True, True, 6),
  )
  def testLayerNormVariables(self, lstm_class, dim, use_bias, use_layer_norm,
                             expected_num_variables):
    batch_size = 2
    input_shape = (8,) * dim + (3,)

    lstm = lstm_class(
        input_shape=input_shape,
        output_channels=5,
        kernel_shape=3,
        use_bias=use_bias,
        use_layer_norm=use_layer_norm)
    inputs = tf.ones(dtype=tf.float32, shape=((1, batch_size) + input_shape))
    initial_state = lstm.initial_state(batch_size, tf.float32)
    tf.nn.dynamic_rnn(lstm,
                      inputs,
                      time_major=True,
                      initial_state=initial_state,
                      dtype=tf.float32)

    self.assertEqual(use_layer_norm, lstm.use_layer_norm)

    # Expect the following variables:
    # Weight, and bias if present, to apply to input
    # Weight, and bias if present, to apply to hidden state
    # LayerNorm's gamma and beta, if present
    self.assertLen(lstm.get_variables(), expected_num_variables)

  @parameterized.parameters(
      (snt.Conv1DLSTM, 1, False),
      (snt.Conv1DLSTM, 1, True),
      (snt.Conv2DLSTM, 2, False),
      (snt.Conv2DLSTM, 2, True),
  )
  def testLayerNorm(self, lstm_class, dim, use_bias):
    """Test that training works, with or without dilated convolutions."""
    time_steps = 3
    batch_size = 2
    input_shape = (8,) * dim
    input_channels = 3
    output_channels = 5
    kernel_shape = 3

    input_shape = (batch_size,) + input_shape + (input_channels,)

    lstm = lstm_class(
        input_shape=input_shape[1:],
        output_channels=output_channels,
        kernel_shape=kernel_shape,
        use_bias=use_bias,
        use_layer_norm=True)
    inputs = tf.random_normal((time_steps,) + input_shape, dtype=tf.float32)
    initial_state = lstm.initial_state(batch_size, tf.float32)

    def loss_fn():
      output, _ = tf.nn.dynamic_rnn(lstm,
                                    inputs,
                                    time_major=True,
                                    initial_state=initial_state,
                                    dtype=tf.float32)

      return tf.reduce_mean(tf.square(output))

    train_op = tf.train.GradientDescentOptimizer(1).minimize(
        loss_fn if tf.executing_eagerly() else loss_fn())
    init = tf.global_variables_initializer()
    self.evaluate(init)
    self.evaluate(train_op)


@contrib_eager.run_all_tests_in_graph_and_eager_modes
class GRUTest(tf.test.TestCase, parameterized.TestCase):

  def testShape(self):
    batch_size = 2
    hidden_size = 4
    inputs = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    state = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    gru = snt.GRU(hidden_size, name="rnn")
    output, next_state = gru(inputs, state)
    shape = np.ndarray((batch_size, hidden_size))
    self.assertShapeEqual(shape, next_state)
    self.assertShapeEqual(shape, output)

  @parameterized.parameters(
      {"input_size": 4, "hidden_size": 5},
      {"input_size": 16, "hidden_size": 128})
  def testVariables(self, input_size, hidden_size):
    batch_size = 5
    mod_name = "rnn"
    inputs = tf.ones(dtype=tf.float32, shape=[batch_size, input_size])
    state = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    gru = snt.GRU(hidden_size, name=mod_name)
    self.assertEqual(gru.scope_name, mod_name)
    with self.assertRaisesRegexp(snt.Error, "not instantiated yet"):
      gru.get_variables()
    gru(inputs, state)

    gru_variables = gru.get_variables()
    self.assertLen(gru_variables, 9, "GRU should have 9 variables")
    param_map = {param.name.split("/")[-1].split(":")[0]: param
                 for param in gru_variables}
    for part in ["z", "r", "h"]:
      self.assertShapeEqual(
          np.ndarray(hidden_size), tf.convert_to_tensor(param_map["b" + part]))
      self.assertShapeEqual(
          np.ndarray((hidden_size, hidden_size)),
          tf.convert_to_tensor(param_map["u" + part]))
      self.assertShapeEqual(
          np.ndarray((input_size, hidden_size)),
          tf.convert_to_tensor(param_map["w" + part]))

  def testComputation(self):

    def sigmoid(x):
      return 1 / (1 + np.exp(-x))

    batch_size = 2
    input_size = 3
    hidden_size = 5

    # With random data, check the TF calculation matches the Numpy version.
    input_data = np.random.randn(batch_size, input_size)
    state_data = np.random.randn(batch_size, hidden_size)
    inputs = tf.constant(input_data)
    state_in = tf.constant(state_data)

    gru = snt.GRU(hidden_size, name="rnn")
    _, state = gru(inputs, state_in)
    gru_variables = gru.get_variables()
    param_map = {param.name.split("/")[-1].split(":")[0]: param
                 for param in gru_variables}

    self.evaluate(tf.global_variables_initializer())
    fetches = [
        state, param_map["wz"], param_map["uz"], param_map["bz"],
        param_map["wr"], param_map["ur"], param_map["br"], param_map["wh"],
        param_map["uh"], param_map["bh"]
    ]
    output = self.evaluate(fetches)

    state_ex, wz, uz, bz, wr, ur, br, wh, uh, bh = output
    z = sigmoid(np.dot(input_data, wz) + np.dot(state_data, uz) + bz)
    r = sigmoid(np.dot(input_data, wr) + np.dot(state_data, ur) + br)
    reset_state = r * state_data
    h_twiddle = np.tanh(np.dot(input_data, wh) + np.dot(reset_state, uh)+ bh)

    state_real = (1 - z) * state_data + z * h_twiddle

    self.assertAllClose(state_real, state_ex)

  def testInitializers(self):
    batch_size = 2
    hidden_size = 4

    # Test we can successfully create the GRU with initializers.
    keys = snt.GRU.POSSIBLE_INITIALIZER_KEYS
    initializers = {
        key: tf.constant_initializer(i) for i, key in enumerate(keys)
    }
    gru = snt.GRU(hidden_size, initializers=initializers)

    # Test we can build the GRU.
    inputs = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    state = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    gru(inputs, state)
    init = tf.global_variables_initializer()

    # Test that the initializers have been correctly applied.
    gru_variables = [getattr(gru, "_" + key) for key in keys]
    self.evaluate(init)
    gru_variables_v = self.evaluate(gru_variables)
    for i, gru_variable_v in enumerate(gru_variables_v):
      self.assertAllClose(gru_variable_v, i * np.ones(gru_variable_v.shape))

  def testPartitioners(self):
    if tf.executing_eagerly():
      self.skipTest("Partitioned variables are not supported in eager mode.")
    batch_size = 2
    hidden_size = 4

    # Test we can successfully create the GRU with partitioners.
    keys = snt.GRU.POSSIBLE_INITIALIZER_KEYS
    partitioners = {
        key: tf.variable_axis_size_partitioner(10) for key in keys
    }
    gru = snt.GRU(hidden_size, partitioners=partitioners)

    # Test we can build the GRU.
    inputs = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    state = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    gru(inputs, state)

    # Test that the variables are partitioned.
    for key in keys:
      self.assertEqual(type(getattr(gru, "_" + key)),
                       variables.PartitionedVariable)

  def testRegularizers(self):
    batch_size = 2
    hidden_size = 4

    # Test we can successfully create the GRU with regularizers.
    keys = snt.GRU.POSSIBLE_INITIALIZER_KEYS
    regularizers = {
        key: tf.nn.l2_loss for key in keys
    }
    gru = snt.GRU(hidden_size, regularizers=regularizers)

    # Test we can build the GRU.
    inputs = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    state = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    gru(inputs, state)

    # Test that we have regularization losses.
    self.assertLen(tf.get_collection(
        tf.GraphKeys.REGULARIZATION_LOSSES), len(keys))


@contrib_eager.run_all_tests_in_graph_and_eager_modes
class HighwayCoreTest(tf.test.TestCase, parameterized.TestCase):

  def testShape(self):
    batch_size = 2
    hidden_size = 4
    inputs = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    state = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    core = snt.HighwayCore(hidden_size, num_layers=3)
    output, next_state = core(inputs, state)
    shape = np.ndarray((batch_size, hidden_size))
    self.assertShapeEqual(shape, next_state)
    self.assertShapeEqual(shape, output)

  def testVariables(self):
    batch_size = 5
    input_size = 10
    hidden_size = 20
    num_layers = 3
    mod_name = "rnn"
    inputs = tf.ones(dtype=tf.float32, shape=[batch_size, input_size])
    state = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    core = snt.HighwayCore(hidden_size, num_layers, name=mod_name)
    self.assertEqual(core.scope_name, mod_name)
    with self.assertRaisesRegexp(snt.Error, "not instantiated yet"):
      core.get_variables()
    core(inputs, state)

    core_variables = core.get_variables()
    self.assertLen(core_variables, 2 + 4 * num_layers)
    param_map = {param.name.split("/")[-1].split(":")[0]: param
                 for param in core_variables}
    self.assertShapeEqual(
        np.ndarray((input_size, hidden_size)),
        tf.convert_to_tensor(param_map["wt"]))
    self.assertShapeEqual(
        np.ndarray((input_size, hidden_size)),
        tf.convert_to_tensor(param_map["wh"]))
    for layer_index in xrange(num_layers):
      layer_str = str(layer_index)
      self.assertShapeEqual(
          np.ndarray(hidden_size),
          tf.convert_to_tensor(param_map["bt" + layer_str]))
      self.assertShapeEqual(
          np.ndarray(hidden_size),
          tf.convert_to_tensor(param_map["bh" + layer_str]))
      self.assertShapeEqual(
          np.ndarray((hidden_size, hidden_size)),
          tf.convert_to_tensor(param_map["wt" + layer_str]))
      self.assertShapeEqual(
          np.ndarray((hidden_size, hidden_size)),
          tf.convert_to_tensor(param_map["wh" + layer_str]))

  @parameterized.parameters(True, False)
  def testComputation(self, with_dropout):
    """Checks that the TF and numpy versions match on random data."""

    def sigmoid(x):
      return 1 / (1 + np.exp(-x))

    batch_size = 2
    input_size = 3
    hidden_size = 5
    num_layers = 2
    input_data = np.random.randn(batch_size, input_size)
    state_data = np.random.randn(batch_size, hidden_size)
    inputs = tf.constant(input_data)
    state_in = tf.constant(state_data)
    if with_dropout:
      core, test_core = snt.highway_core_with_recurrent_dropout(
          hidden_size, num_layers, keep_prob=1.0)
      initial_state = core.initial_state(batch_size, dtype=tf.float64)
      _, state = core(inputs, (state_in, initial_state[1]))
      core_variables = test_core.get_variables()
    else:
      core = snt.HighwayCore(hidden_size, num_layers, name="rnn")
      _, state = core(inputs, state_in)
      core_variables = core.get_variables()
    param_map = {param.name.split("/")[-1].split(":")[0]: param
                 for param in core_variables}

    param_names = ["wt", "wh"]
    param_names += ["wt0", "bt0", "wh0", "bh0", "wt1", "bt1", "wh1", "bh1"]

    self.evaluate(tf.global_variables_initializer())
    fetches = [state] + [param_map[name] for name in param_names]
    output = self.evaluate(fetches)

    state_ex, wt, wh, wt0, bt0, wh0, bh0, wt1, bt1, wh1, bh1 = output
    # Layer 1 computation.
    t = sigmoid(np.dot(input_data, wt) + np.dot(state_data, wt0) + bt0)
    h = np.tanh(np.dot(input_data, wh) + np.dot(state_data, wh0) + bh0)
    state_data = (1 - t) * state_data + t * h
    # Layer 2 computation.
    t = sigmoid(np.dot(state_data, wt1) + bt1)
    h = np.tanh(np.dot(state_data, wh1) + bh1)
    state_data = (1 - t) * state_data + t * h

    if with_dropout:
      state_ex = state_ex[0]
    self.assertAllClose(state_data, state_ex)


@contrib_eager.run_all_tests_in_graph_and_eager_modes
class LSTMBlockCellTest(tf.test.TestCase, parameterized.TestCase):

  def testShape(self):
    batch_size = 2
    hidden_size = 5
    inputs = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    prev_hidden = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    prev_cell = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    lstm = snt.LSTMBlockCell(hidden_size)
    output, next_state = lstm(inputs, (prev_hidden, prev_cell))

    shape = np.ndarray((batch_size, hidden_size))

    self.assertShapeEqual(shape, next_state[0])
    self.assertShapeEqual(shape, next_state[1])
    self.assertShapeEqual(shape, output)

    self.assertEqual(hidden_size, lstm.output_size)
    self.assertEqual((hidden_size, hidden_size), lstm.state_size)

  def testVariables(self):
    batch_size = 5
    hidden_size = 20
    mod_name = "lstm_block"
    inputs = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    prev_cell = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    prev_hidden = tf.ones(dtype=tf.float32, shape=[batch_size, hidden_size])
    lstm = snt.LSTMBlockCell(hidden_size, name=mod_name)
    self.assertEqual(lstm.scope_name, mod_name)
    with self.assertRaisesRegexp(snt.Error, "not instantiated yet"):
      lstm.get_variables()
    lstm(inputs, (prev_hidden, prev_cell))

    lstm_variables = lstm.get_variables()
    self.assertLen(lstm_variables, 2, "LSTM should have 2 variables")
    param_map = {param.name.split("/")[-1].split(":")[0]:
                 param for param in lstm_variables}

    self.assertShapeEqual(np.ndarray(4 * hidden_size),
                          tf.convert_to_tensor(param_map["bias"]))
    self.assertShapeEqual(np.ndarray((2 * hidden_size, 4 * hidden_size)),
                          tf.convert_to_tensor(param_map["kernel"]))


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