# 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 layers."""

import lingvo.compat as tf
from lingvo.core import batch_major_attention
from lingvo.core import py_utils
from lingvo.core import test_utils
from lingvo.tasks.lm import layers as lm_layers
import numpy as np

FLAGS = tf.flags.FLAGS


class RnnLmNoEmbeddingTest(test_utils.TestCase):

  def _testParams(self, dims, vocab):
    p = lm_layers.RnnLmNoEmbedding.Params()
    p.name = 'rnnlm'
    p.vocab_size = vocab
    p.rnns.cell_tpl.num_output_nodes = dims
    p.rnns.cell_tpl.num_input_nodes = dims
    p.softmax.input_dim = dims
    p.softmax.num_classes = vocab
    return p

  def testBasic(self):
    time, batch, dims, vocab = 5, 3, 6, 8
    p = self._testParams(dims, vocab)

    with self.session(use_gpu=True):
      lm = p.Instantiate()
      np.random.seed(12345)
      inputs = np.random.normal(size=[time, batch, dims])
      inputs = tf.constant(inputs, tf.float32)
      paddings = np.zeros([time, batch])
      paddings[-1] = 1.0
      paddings = tf.constant(paddings, tf.float32)
      targets = tf.constant(
          np.random.randint(vocab, size=(time, batch)), tf.int32)
      self.evaluate(tf.global_variables_initializer())
      xent_output, state1 = lm.FPropDefaultTheta(
          inputs=inputs,
          paddings=paddings,
          state0=lm.zero_state(lm.theta, batch),
          labels=py_utils.NestedMap(
              class_weights=1 - paddings, class_ids=targets))

      xent_output_val, state1_val = self.evaluate([xent_output, state1])

      print('xent_output_val', xent_output_val)
      print('state1', state1_val)
      test_utils.CompareToGoldenSingleFloat(self, 2.1042602, xent_output_val.avg_xent)  # pyformat: disable pylint: disable=line-too-long
      self.assertAllEqual(xent_output_val.per_example_argmax,
                          np.argmax(xent_output_val.logits, axis=-1))

  def testProjection(self):
    time, batch, dims, proj, vocab = 5, 3, 6, 4, 8
    p = self._testParams(dims, vocab)
    p.rnns.cell_tpl.num_output_nodes = proj
    p.rnns.cell_tpl.num_hidden_nodes = dims
    p.softmax.input_dim = proj

    with self.session(use_gpu=True):
      lm = p.Instantiate()
      np.random.seed(12345)
      inputs = np.random.normal(size=[time, batch, dims])
      inputs = tf.constant(inputs, tf.float32)
      paddings = np.zeros([time, batch])
      paddings[-1] = 1.0
      paddings = tf.constant(paddings, tf.float32)
      targets = tf.constant(
          np.random.randint(vocab, size=(time, batch)), tf.int32)
      self.evaluate(tf.global_variables_initializer())
      xent_output, state1 = lm.FPropDefaultTheta(
          inputs=inputs,
          paddings=paddings,
          state0=lm.zero_state(lm.theta, batch),
          labels=py_utils.NestedMap(
              class_weights=1 - paddings, class_ids=targets))

      xent_output_val, state1_val = self.evaluate([xent_output, state1])

      print('xent_output_val', xent_output_val)
      print('state1', state1_val)
      test_utils.CompareToGoldenSingleFloat(self, 2.1322777, xent_output_val.avg_xent)  # pyformat: disable pylint: disable=line-too-long
      self.assertAllEqual(xent_output_val.per_example_argmax,
                          np.argmax(xent_output_val.logits, axis=-1))

  def testBasicGrad(self):
    time, batch, dims, vocab = 5, 3, 6, 8
    p = self._testParams(dims, vocab)
    p.dtype = tf.float64

    with self.session(use_gpu=False, graph=tf.Graph()) as sess:
      lm = p.Instantiate()
      np.random.seed(12345)
      inputs = np.random.normal(size=[time, batch, dims])
      inputs = tf.constant(inputs, tf.float64)
      paddings = np.zeros([time, batch])
      paddings[-1] = 1.0
      paddings = tf.constant(paddings, tf.float64)
      targets = tf.constant(
          np.random.randint(vocab, size=(time, batch)), tf.int32)

      xent_output, _ = lm.FPropDefaultTheta(
          inputs=inputs,
          paddings=paddings,
          state0=lm.zero_state(lm.theta, batch),
          labels=py_utils.NestedMap(
              class_weights=1 - paddings, class_ids=targets))

      lm_vars = lm.vars.Flatten()
      # Now add the backward graph.
      grads = tf.gradients(xent_output.avg_xent, lm_vars)

      self.evaluate(tf.global_variables_initializer())
      self.assertEqual(len(lm_vars), len(grads))
      for x, grad_x in zip(lm_vars, grads):
        grad_symbolic = self.evaluate(grad_x)
        grad_numeric = test_utils.ComputeNumericGradient(
            sess, xent_output.avg_xent, x, delta=1e-6)
        self.assertAllClose(grad_symbolic, grad_numeric, atol=0.005)

  def testDirectFeatures(self,):
    time, batch, dims, vocab = 5, 3, 6, 8
    p = self._testParams(dims, vocab)
    direct_features_dim = 4
    p.direct_features_dim = direct_features_dim
    p.softmax.input_dim = dims + direct_features_dim

    with self.session(use_gpu=True):
      lm = p.Instantiate()
      np.random.seed(12345)
      inputs = np.random.normal(size=[time, batch, dims])
      inputs = tf.constant(inputs, tf.float32)

      direct_features = np.random.normal(
          size=[time, batch, direct_features_dim])
      direct_features = tf.constant(direct_features, tf.float32)

      paddings = np.zeros([time, batch])
      paddings[-1] = 1.0
      paddings = tf.constant(paddings, tf.float32)
      targets = tf.constant(
          np.random.randint(vocab, size=(time, batch)), tf.int32)
      self.evaluate(tf.global_variables_initializer())
      xent_output, state1 = lm.FPropDefaultTheta(
          inputs=inputs,
          paddings=paddings,
          state0=lm.zero_state(lm.theta, batch),
          labels=py_utils.NestedMap(
              class_weights=1 - paddings, class_ids=targets),
          direct_features=direct_features)

      xent_output_val, state1_val = self.evaluate([xent_output, state1])

      print('xent_output_val', xent_output_val)
      print('state1', state1_val)
      test_utils.CompareToGoldenSingleFloat(self, 2.3522419929504395, xent_output_val.avg_xent)  # pyformat: disable pylint: disable=line-too-long
      self.assertAllEqual(xent_output_val.per_example_argmax,
                          np.argmax(xent_output_val.logits, axis=-1))

  def testCombineStates(self):
    time, batch, dims, vocab = 5, 3, 6, 8
    p = self._testParams(dims, vocab)

    with self.session(use_gpu=True):
      lm = p.Instantiate()
      np.random.seed(12345)
      inputs = np.random.normal(size=[time, batch, dims])
      inputs = tf.constant(inputs, tf.float32)
      paddings = np.zeros([time, batch])
      paddings[-1] = 1.0
      paddings = tf.constant(paddings, tf.float32)
      targets = tf.constant(
          np.random.randint(vocab, size=(time, batch)), tf.int32)
      self.evaluate(tf.global_variables_initializer())

      state0 = lm.zero_state(lm.theta, batch)
      _, state1 = lm.FPropDefaultTheta(
          inputs=inputs,
          paddings=paddings,
          state0=state0,
          labels=py_utils.NestedMap(
              class_weights=1 - paddings, class_ids=targets))
      switch = tf.constant([True, True, False], dtype=tf.bool)
      combined_state = lm.CombineStates(state0, state1, switch)
      state0_val, state1_val, combined_state_val = self.evaluate(
          [state0, state1, combined_state])

      print('state1_val', state1_val)
      print('combined_state_val', combined_state_val)
      combined_m_state = combined_state_val.rnn[0].m
      combined_c_state = combined_state_val.rnn[0].c
      self.assertAllEqual(combined_m_state[0], state0_val.rnn[0].m[0])
      self.assertAllEqual(combined_c_state[0], state0_val.rnn[0].c[0])
      self.assertAllEqual(combined_m_state[1], state0_val.rnn[0].m[1])
      self.assertAllEqual(combined_c_state[1], state0_val.rnn[0].c[1])
      self.assertAllEqual(combined_m_state[2], state1_val.rnn[0].m[2])
      self.assertAllEqual(combined_c_state[2], state1_val.rnn[0].c[2])


class RnnLmTest(test_utils.TestCase):

  def testBasic(self):
    time, batch, dims, vocab = 5, 3, 6, 8

    p = lm_layers.RnnLm.Params()
    p.name = 'rnnlm'
    p.vocab_size = vocab
    p.emb.vocab_size = vocab
    p.emb.embedding_dim = dims
    p.rnns.cell_tpl.num_output_nodes = dims
    p.rnns.cell_tpl.num_input_nodes = dims
    p.softmax.input_dim = dims
    p.softmax.num_classes = vocab

    with self.session(use_gpu=True):
      lm = p.Instantiate()
      np.random.seed(12345)
      inputs = np.random.randint(vocab, size=[time, batch])
      targets = np.zeros([time, batch])
      targets[:-1] = inputs[1:]
      inputs = tf.constant(inputs, tf.int32)
      paddings = np.zeros([time, batch])
      paddings[-1] = 1.0
      paddings = tf.constant(paddings, tf.float32)
      targets = tf.constant(targets, tf.int32)
      self.evaluate(tf.global_variables_initializer())
      xent_output, state1 = lm.FPropDefaultTheta(
          inputs=inputs,
          paddings=paddings,
          state0=lm.zero_state(lm.theta, batch),
          labels=py_utils.NestedMap(
              class_weights=1 - paddings, class_ids=targets))

      xent_output_val, state1_val = self.evaluate([xent_output, state1])

      print('xent_output_val', xent_output_val)
      print('state1', state1_val)
      test_utils.CompareToGoldenSingleFloat(self, 2.0853612, xent_output_val.avg_xent)  # pyformat: disable pylint: disable=line-too-long
      self.assertAllEqual(xent_output_val.per_example_argmax,
                          np.argmax(xent_output_val.logits, axis=-1))

  def testDropout(self):
    seed = 12345
    tf.random.set_seed(seed)
    np.random.seed(seed)

    time, batch, dims, vocab = 5, 3, 6, 8

    p = lm_layers.RnnLm.Params()
    p.name = 'rnnlm'
    p.vocab_size = vocab
    p.emb.vocab_size = vocab
    p.emb.embedding_dim = dims
    p.embedding_dropout_keep_prob = 0.8
    p.embedding_dropout_seed = seed
    p.rnns.cell_tpl.num_output_nodes = dims
    p.rnns.cell_tpl.num_input_nodes = dims
    p.softmax.input_dim = dims
    p.softmax.num_classes = vocab

    with self.session(use_gpu=True):
      lm = p.Instantiate()
      inputs = np.random.randint(vocab, size=[time, batch])
      targets = np.zeros([time, batch])
      targets[:-1] = inputs[1:]
      inputs = tf.constant(inputs, tf.int32)
      paddings = np.zeros([time, batch])
      paddings[-1] = 1.0
      paddings = tf.constant(paddings, tf.float32)
      targets = tf.constant(targets, tf.int32)
      self.evaluate(tf.global_variables_initializer())
      xent_output, state1 = lm.FPropDefaultTheta(
          inputs=inputs,
          paddings=paddings,
          state0=lm.zero_state(lm.theta, batch),
          labels=py_utils.NestedMap(
              class_weights=1 - paddings, class_ids=targets))

      xent_output_val, state1_val = self.evaluate([xent_output, state1])

      print('xent_output_val', xent_output_val)
      print('state1', state1_val)
      test_utils.CompareToGoldenSingleFloat(self, 2.084798, xent_output_val.avg_xent)  # pyformat: disable pylint: disable=line-too-long
      self.assertAllEqual(xent_output_val.per_example_argmax,
                          np.argmax(xent_output_val.logits, axis=-1))

  def testBasicGrad(self):
    time, batch, dims, vocab = 5, 3, 6, 8

    p = lm_layers.RnnLm.Params()
    p.dtype = tf.float64
    p.name = 'rnnlm'
    p.vocab_size = vocab
    p.emb.vocab_size = vocab
    p.emb.embedding_dim = dims
    p.rnns.cell_tpl.num_output_nodes = dims
    p.rnns.cell_tpl.num_input_nodes = dims
    p.softmax.input_dim = dims
    p.softmax.num_classes = vocab

    with self.session(use_gpu=False, graph=tf.Graph()) as sess:
      lm = p.Instantiate()
      np.random.seed(12345)
      inputs = np.random.randint(vocab, size=[time, batch])
      targets = np.zeros([time, batch])
      targets[:-1] = inputs[1:]
      inputs = tf.constant(inputs, tf.int32)
      paddings = np.zeros([time, batch])
      paddings[-1] = 1.0
      paddings = tf.constant(paddings, tf.float64)
      targets = tf.constant(targets, tf.int32)
      self.evaluate(tf.global_variables_initializer())
      xent_output, _ = lm.FPropDefaultTheta(
          inputs=inputs,
          paddings=paddings,
          state0=lm.zero_state(lm.theta, batch),
          labels=py_utils.NestedMap(
              class_weights=1 - paddings, class_ids=targets))

      lm_vars = lm.vars.Flatten()
      # Now add the backward graph.
      grads = tf.gradients(xent_output.avg_xent, lm_vars)

      for i, x in enumerate(grads):
        if isinstance(x, tf.IndexedSlices):
          grads[i] = tf.math.unsorted_segment_sum(x.values, x.indices,
                                                  x.dense_shape[0])

      self.evaluate(tf.global_variables_initializer())
      self.assertEqual(len(lm_vars), len(grads))
      for x, grad_x in zip(lm_vars, grads):
        grad_symbolic = self.evaluate(grad_x)
        grad_numeric = test_utils.ComputeNumericGradient(
            sess, xent_output.avg_xent, x, delta=1e-6)
        self.assertAllClose(grad_symbolic, grad_numeric, atol=0.005)


class ConditionalRnnLmTest(test_utils.TestCase):

  def testBasic(self):
    time, batch, dims, vocab, condition_dim = 5, 3, 6, 8, 7

    p = lm_layers.ConditionalRnnLm.Params()
    p.name = 'conditionalrnnlm'
    p.vocab_size = vocab
    p.emb.vocab_size = vocab
    p.emb.embedding_dim = dims
    model_dim = dims + condition_dim
    p.rnns.cell_tpl.num_output_nodes = model_dim
    p.rnns.cell_tpl.num_input_nodes = model_dim
    p.softmax.input_dim = model_dim
    p.softmax.num_classes = vocab
    p.condition_dim = condition_dim

    with self.session(use_gpu=True):
      lm = p.Instantiate()
      np.random.seed(12345)
      inputs = np.random.randint(vocab, size=[time, batch])
      targets = np.zeros([time, batch])
      targets[:-1] = inputs[1:]
      inputs = tf.constant(inputs, tf.int32)
      paddings = np.zeros([time, batch])
      paddings[-1] = 1.0
      paddings = tf.constant(paddings, tf.float32)
      targets = tf.constant(targets, tf.int32)
      condition = tf.constant(np.ones([batch, condition_dim]), tf.float64)
      self.evaluate(tf.global_variables_initializer())
      xent_output, state1 = lm.FPropDefaultTheta(
          inputs=inputs,
          paddings=paddings,
          state0=lm.zero_state(lm.theta, batch),
          condition=condition,
          labels=py_utils.NestedMap(
              class_weights=1 - paddings, class_ids=targets))

      xent_output_val, state1_val = self.evaluate([xent_output, state1])

      print('xent_output_val', xent_output_val)
      print('state1', state1_val)
      test_utils.CompareToGoldenSingleFloat(self, 2.10713076, xent_output_val.avg_xent)  # pyformat: disable pylint: disable=line-too-long
      self.assertAllEqual(xent_output_val.per_example_argmax,
                          np.argmax(xent_output_val.logits, axis=-1))

  def testDropout(self):
    seed = 12345
    tf.random.set_seed(seed)
    np.random.seed(seed)

    time, batch, dims, vocab, condition_dim = 5, 3, 6, 8, 7

    p = lm_layers.ConditionalRnnLm.Params()
    p.name = 'conditionalrnnlm'
    p.vocab_size = vocab
    p.emb.vocab_size = vocab
    p.emb.embedding_dim = dims
    p.embedding_dropout_keep_prob = 0.8
    p.embedding_dropout_seed = seed
    model_dim = dims + condition_dim
    p.rnns.cell_tpl.num_output_nodes = model_dim
    p.rnns.cell_tpl.num_input_nodes = model_dim
    p.softmax.input_dim = model_dim
    p.softmax.num_classes = vocab
    p.condition_dim = condition_dim

    with self.session(use_gpu=True):
      lm = p.Instantiate()
      inputs = np.random.randint(vocab, size=[time, batch])
      targets = np.zeros([time, batch])
      targets[:-1] = inputs[1:]
      inputs = tf.constant(inputs, tf.int32)
      paddings = np.zeros([time, batch])
      paddings[-1] = 1.0
      paddings = tf.constant(paddings, tf.float32)
      targets = tf.constant(targets, tf.int32)
      condition = tf.constant(np.ones([batch, condition_dim]), tf.float64)
      self.evaluate(tf.global_variables_initializer())
      xent_output, state1 = lm.FPropDefaultTheta(
          inputs=inputs,
          paddings=paddings,
          state0=lm.zero_state(lm.theta, batch),
          condition=condition,
          labels=py_utils.NestedMap(
              class_weights=1 - paddings, class_ids=targets))

      xent_output_val, state1_val = self.evaluate([xent_output, state1])

      print('xent_output_val', xent_output_val)
      print('state1', state1_val)
      test_utils.CompareToGoldenSingleFloat(self, 2.17278885, xent_output_val.avg_xent)  # pyformat: disable pylint: disable=line-too-long
      self.assertAllEqual(xent_output_val.per_example_argmax,
                          np.argmax(xent_output_val.logits, axis=-1))

  def testBasicGrad(self):
    time, batch, dims, vocab, condition_dim = 5, 3, 6, 8, 7

    p = lm_layers.ConditionalRnnLm.Params()
    p.name = 'conditionalrnnlm'
    p.dtype = tf.float64
    p.vocab_size = vocab
    p.emb.vocab_size = vocab
    p.emb.embedding_dim = dims
    model_dim = dims + condition_dim
    p.rnns.cell_tpl.num_output_nodes = model_dim
    p.rnns.cell_tpl.num_input_nodes = model_dim
    p.softmax.input_dim = model_dim
    p.softmax.num_classes = vocab
    p.condition_dim = condition_dim

    with self.session(use_gpu=False, graph=tf.Graph()) as sess:
      lm = p.Instantiate()
      np.random.seed(12345)
      inputs = np.random.randint(vocab, size=[time, batch])
      targets = np.zeros([time, batch])
      targets[:-1] = inputs[1:]
      inputs = tf.constant(inputs, tf.int32)
      paddings = np.zeros([time, batch])
      paddings[-1] = 1.0
      paddings = tf.constant(paddings, tf.float64)
      targets = tf.constant(targets, tf.int32)
      condition = tf.constant(np.ones([batch, condition_dim]), tf.float64)
      self.evaluate(tf.global_variables_initializer())
      xent_output, _ = lm.FPropDefaultTheta(
          inputs=inputs,
          paddings=paddings,
          state0=lm.zero_state(lm.theta, batch),
          condition=condition,
          labels=py_utils.NestedMap(
              class_weights=1 - paddings, class_ids=targets))

      lm_vars = lm.vars.Flatten()
      # Now add the backward graph.
      grads = tf.gradients(xent_output.avg_xent, lm_vars)

      for i, x in enumerate(grads):
        if isinstance(x, tf.IndexedSlices):
          grads[i] = tf.math.unsorted_segment_sum(x.values, x.indices,
                                                  x.dense_shape[0])

      self.evaluate(tf.global_variables_initializer())
      self.assertEqual(len(lm_vars), len(grads))
      for x, grad_x in zip(lm_vars, grads):
        grad_symbolic = self.evaluate(grad_x)
        grad_numeric = test_utils.ComputeNumericGradient(
            sess, xent_output.avg_xent, x, delta=1e-6)
        self.assertAllClose(grad_symbolic, grad_numeric, atol=0.005)


class MoeLmTest(test_utils.TestCase):

  def _MoeLmParams(self, vocab, shared_emb, add_postgating_rnn=True):
    p = lm_layers.MoeLm.Params()
    p.name = 'moelm'
    p.vocab_size = vocab
    p.emb.vocab_size = vocab
    p.emb.embedding_dim = 17
    p.shared_emb = shared_emb
    p.add_postgating_rnn = add_postgating_rnn
    p.rnns.cell_tpl.num_input_nodes = 17
    p.rnns.cell_tpl.num_output_nodes = 23
    p.number_of_experts = 4
    p.merge.vocab_size = vocab
    p.merge.rnns.cell_tpl.num_input_nodes = 23
    p.merge.rnns.cell_tpl.num_output_nodes = 32
    p.merge.softmax.input_dim = 32
    p.merge.softmax.num_classes = vocab
    return p

  def _GetData(self, vocab, time, batch):
    inputs = np.random.randint(vocab, size=[time, batch])
    targets = np.zeros([time, batch])
    targets[:-1] = inputs[1:]
    inputs = tf.constant(inputs, tf.int32)
    paddings = np.zeros([time, batch])
    paddings[-1] = 1.0
    paddings = tf.constant(paddings, tf.float32)
    targets = tf.constant(targets, tf.int32)
    labels = py_utils.NestedMap(class_weights=1 - paddings, class_ids=targets)
    return inputs, paddings, labels

  def testFPropNoSharedEmb(self):
    vocab, time, batch = 7, 13, 3
    p = self._MoeLmParams(vocab, False)

    with self.session(graph=tf.Graph()):
      np.random.seed(54321)
      tf.random.set_seed(123456)
      lm = p.Instantiate()
      inputs, paddings, labels = self._GetData(vocab, time, batch)
      self.evaluate(tf.global_variables_initializer())
      xent_output, state1 = lm.FPropDefaultTheta(
          inputs=inputs,
          paddings=paddings,
          state0=lm.zero_state(lm.theta, batch),
          labels=labels)

      xent_output_val, state1_val = self.evaluate([xent_output, state1])

      print('xent_output_val', xent_output_val)
      print('state1', state1_val)
      test_utils.CompareToGoldenSingleFloat(self, 1.9460623, xent_output_val.avg_xent)  # pyformat: disable pylint: disable=line-too-long
      self.assertAllEqual(xent_output_val.per_example_argmax,
                          np.argmax(xent_output_val.logits, axis=-1))

  def testFPropSharedEmb(self):
    vocab, time, batch = 7, 13, 3
    p = self._MoeLmParams(vocab, False)

    with self.session(graph=tf.Graph()):
      np.random.seed(54321)
      tf.random.set_seed(123456)
      lm = p.Instantiate()
      inputs, paddings, labels = self._GetData(vocab, time, batch)
      self.evaluate(tf.global_variables_initializer())
      xent_output, state1 = lm.FPropDefaultTheta(
          inputs=inputs,
          paddings=paddings,
          state0=lm.zero_state(lm.theta, batch),
          labels=labels)

      xent_output_val, state1_val = self.evaluate([xent_output, state1])

      print('xent_output_val', xent_output_val)
      print('state1', state1_val)
      test_utils.CompareToGoldenSingleFloat(self, 1.9460623, xent_output_val.avg_xent)  # pyformat: disable pylint: disable=line-too-long
      self.assertAllEqual(xent_output_val.per_example_argmax,
                          np.argmax(xent_output_val.logits, axis=-1))

  def testFPropNoPostGatingRNN(self):
    vocab, time, batch = 7, 13, 3
    p = self._MoeLmParams(vocab, False, False)

    with self.session(graph=tf.Graph()):
      np.random.seed(54321)
      tf.random.set_seed(123456)
      lm = p.Instantiate()
      inputs, paddings, labels = self._GetData(vocab, time, batch)
      self.evaluate(tf.global_variables_initializer())
      xent_output, state1 = lm.FPropDefaultTheta(
          inputs=inputs,
          paddings=paddings,
          state0=lm.zero_state(lm.theta, batch),
          labels=labels)

      xent_output_val, state1_val = self.evaluate([xent_output, state1])

      print('xent_output_val', xent_output_val)
      print('state1', state1_val)
      test_utils.CompareToGoldenSingleFloat(self, 1.9443978, xent_output_val.avg_xent)  # pyformat: disable pylint: disable=line-too-long
      self.assertAllEqual(xent_output_val.per_example_argmax,
                          np.argmax(xent_output_val.logits, axis=-1))

  def testBProp(self):
    vocab, time, batch = 7, 4, 3
    p = self._MoeLmParams(vocab, True)
    p.dtype = tf.float64

    with self.session(graph=tf.Graph()) as sess:
      np.random.seed(54321)
      tf.random.set_seed(123456)
      lm = p.Instantiate()
      inputs, paddings, labels = self._GetData(vocab, time, batch)
      self.evaluate(tf.global_variables_initializer())
      xent_output, _ = lm.FPropDefaultTheta(
          inputs=inputs,
          paddings=tf.cast(paddings, p.dtype),
          state0=lm.zero_state(lm.theta, batch),
          labels=labels)

      lm_vars = lm.vars.Flatten()
      # Now add the backward graph.
      grads = tf.gradients(xent_output.avg_xent, lm_vars)

      for i, x in enumerate(grads):
        if isinstance(x, tf.IndexedSlices):
          grads[i] = tf.math.unsorted_segment_sum(x.values, x.indices,
                                                  x.dense_shape[0])

      self.evaluate(tf.global_variables_initializer())
      self.assertEqual(len(lm_vars), len(grads))
      step = 11  # Speed up the test.
      for x, grad_x in zip(lm_vars, grads):
        grad_symbolic = self.evaluate(grad_x)
        grad_numeric = test_utils.ComputeNumericGradient(
            sess, xent_output.avg_xent, x, step=step, delta=1e-6)
        self.assertAllClose(
            grad_symbolic.reshape([-1])[::step],
            grad_numeric.reshape([-1])[::step])


class TransformerLmNoEmbeddingTest(test_utils.TestCase):

  def _testParams(self, dtype=tf.float32):
    model_dim, hidden_dim, vocab_size = 4, 6, 8
    p = lm_layers.TransformerLmNoEmbedding.Params()
    p.name = 'xformerlm'
    p.random_seed = 93820986
    p.dtype = dtype
    p.vocab_size = vocab_size
    p.model_dim = model_dim
    p.num_trans_layers = 3
    p.position_emb.embedding_dim = model_dim
    p.trans_tpl.source_dim = model_dim
    p.trans_tpl.tr_atten_tpl.num_attention_heads = 2
    p.trans_tpl.tr_fflayer_tpl.hidden_dim = hidden_dim
    p.softmax.input_dim = model_dim
    p.softmax.num_classes = vocab_size
    return p

  def _testInputs(self, dtype=tf.float32, last_padding=1.0):
    time, batch, model_dim, vocab_size = 5, 3, 4, 8
    np.random.seed(12345)
    inputs = np.random.normal(size=[time, batch, model_dim])
    inputs = tf.constant(inputs, dtype)
    paddings = np.zeros([time, batch])
    paddings[-1] = last_padding
    paddings = tf.constant(paddings, dtype)
    targets = tf.constant(
        np.random.randint(vocab_size, size=(time, batch)), tf.int32)
    return inputs, paddings, targets

  def testBasic(self):
    p = self._testParams(dtype=tf.float32)
    with self.session(use_gpu=True):
      lm = p.Instantiate()
      inputs, paddings, targets = self._testInputs(dtype=tf.float32)
      self.evaluate(tf.global_variables_initializer())
      xent_output, _ = lm.FPropDefaultTheta(
          inputs=inputs,
          paddings=paddings,
          labels=py_utils.NestedMap(
              class_weights=1 - paddings, class_ids=targets))
      xent_output_val = self.evaluate(xent_output)
      print('xformer xent_output_val.avg_xent', xent_output_val.avg_xent)
      test_utils.CompareToGoldenSingleFloat(self, 3.018613, xent_output_val.avg_xent)  # pyformat: disable pylint: disable=line-too-long
      self.assertAllEqual(xent_output_val.per_example_argmax,
                          np.argmax(xent_output_val.logits, axis=-1))

  def testBasicGrad(self):
    p = self._testParams(dtype=tf.float64)
    with self.session(use_gpu=False, graph=tf.Graph()) as sess:
      lm = p.Instantiate()
      inputs, paddings, targets = self._testInputs(dtype=tf.float64)
      xent_output, _ = lm.FPropDefaultTheta(
          inputs=inputs,
          paddings=paddings,
          labels=py_utils.NestedMap(
              class_weights=1 - paddings, class_ids=targets))

      lm_vars = lm.vars.Flatten()
      # Now add the backward graph.
      grads = tf.gradients(xent_output.avg_xent, lm_vars)

      self.evaluate(tf.global_variables_initializer())
      self.assertEqual(len(lm_vars), len(grads))
      for x, grad_x in zip(lm_vars, grads):
        grad_symbolic = self.evaluate(grad_x)
        grad_numeric = test_utils.ComputeNumericGradient(
            sess, xent_output.avg_xent, x, delta=1e-6)
        self.assertAllClose(grad_symbolic, grad_numeric, atol=0.005)

  def testStep(self):
    p = self._testParams(dtype=tf.float32)
    with self.session(use_gpu=True):
      lm = p.Instantiate()
      inputs, paddings, _ = self._testInputs(dtype=tf.float32, last_padding=0.0)

      self.evaluate(tf.global_variables_initializer())
      xent_output, _ = lm.FPropDefaultTheta(inputs=inputs, paddings=paddings)
      logits1 = xent_output.logits

      time, batch = 5, 3
      prefix_states = lm.zero_state(lm.theta, batch)
      logits2 = []
      for i in range(time):
        l_i_out, prefix_states = lm.Step(lm.theta, inputs[i, :, :],
                                         paddings[i, :], prefix_states)
        logits2.append(l_i_out.logits)
      logits2 = tf.stack(logits2)

      self.evaluate(tf.global_variables_initializer())
      logits1_v, logits2_v = self.evaluate([logits1, logits2])
      print('xformer logits1_v', logits1_v)
      print('xformer logits2_v', logits2_v)
      self.assertAllClose(logits1_v, logits2_v, 1e-5, 1e-5)


class TransformerLmTest(test_utils.TestCase):

  def testBasic(self):
    time, batch, dims, hidden_dim, vocab = 5, 3, 6, 4, 8

    p = lm_layers.TransformerLm.Params()
    p.name = 'transformerlm'
    p.vocab_size = vocab
    p.emb.vocab_size = vocab
    p.emb.embedding_dim = dims
    p.model_dim = dims
    p.num_trans_layers = 3
    p.position_emb.embedding_dim = dims
    p.trans_tpl.source_dim = dims
    p.trans_tpl.tr_atten_tpl.num_attention_heads = 2
    p.trans_tpl.tr_fflayer_tpl.hidden_dim = hidden_dim
    p.softmax.input_dim = dims
    p.softmax.num_classes = vocab

    with self.session(use_gpu=True):
      lm = p.Instantiate()
      np.random.seed(12345)
      inputs = np.random.randint(vocab, size=[time, batch])
      targets = np.zeros([time, batch])
      targets[:-1] = inputs[1:]
      inputs = tf.constant(inputs, tf.int32)
      paddings = np.zeros([time, batch])
      paddings[-1] = 1.0
      paddings = tf.constant(paddings, tf.float32)
      targets = tf.constant(targets, tf.int32)
      self.evaluate(tf.global_variables_initializer())
      xent_output, _ = lm.FPropDefaultTheta(
          inputs=inputs,
          paddings=paddings,
          labels=py_utils.NestedMap(
              class_weights=1 - paddings, class_ids=targets))

      xent_output_val = self.evaluate(xent_output)

      print('xent_output_val', xent_output_val)
      test_utils.CompareToGoldenSingleFloat(self, 3.0489848, xent_output_val.avg_xent)  # pyformat: disable pylint: disable=line-too-long
      self.assertAllEqual(xent_output_val.per_example_argmax,
                          np.argmax(xent_output_val.logits, axis=-1))

  def testDropout(self):
    seed = 12345
    tf.random.set_seed(seed)
    np.random.seed(seed)

    time, batch, dims, hidden_dim, vocab = 5, 3, 6, 4, 8

    p = lm_layers.TransformerLm.Params()
    p.name = 'transformerlm'
    p.vocab_size = vocab
    p.emb.vocab_size = vocab
    p.emb.embedding_dim = dims
    p.model_dim = dims
    p.num_trans_layers = 3
    p.trans_tpl.source_dim = dims
    p.trans_tpl.tr_atten_tpl.num_attention_heads = 2
    p.trans_tpl.tr_fflayer_tpl.hidden_dim = hidden_dim
    p.softmax.input_dim = dims
    p.softmax.num_classes = vocab

    with self.session(use_gpu=True):
      lm = p.Instantiate()
      inputs = np.random.randint(vocab, size=[time, batch])
      targets = np.zeros([time, batch])
      targets[:-1] = inputs[1:]
      inputs = tf.constant(inputs, tf.int32)
      paddings = np.zeros([time, batch])
      paddings[-1] = 1.0
      paddings = tf.constant(paddings, tf.float32)
      targets = tf.constant(targets, tf.int32)
      self.evaluate(tf.global_variables_initializer())
      xent_output, _ = lm.FPropDefaultTheta(
          inputs=inputs,
          paddings=paddings,
          labels=py_utils.NestedMap(
              class_weights=1 - paddings, class_ids=targets))

      xent_output_val = self.evaluate(xent_output)

      print('xent_output_val', xent_output_val)
      test_utils.CompareToGoldenSingleFloat(self, 3.038596, xent_output_val.avg_xent)  # pyformat: disable pylint: disable=line-too-long
      self.assertAllEqual(xent_output_val.per_example_argmax,
                          np.argmax(xent_output_val.logits, axis=-1))

  def testBasicGrad(self):
    time, batch, dims, hidden_dim, vocab = 5, 3, 6, 4, 8

    p = lm_layers.TransformerLm.Params()
    p.dtype = tf.float64
    p.name = 'transformerlm'
    p.vocab_size = vocab
    p.emb.vocab_size = vocab
    p.emb.embedding_dim = dims
    p.model_dim = dims
    p.num_trans_layers = 1
    p.trans_tpl.source_dim = dims
    p.trans_tpl.tr_atten_tpl.num_attention_heads = 2
    p.trans_tpl.tr_fflayer_tpl.hidden_dim = hidden_dim
    p.softmax.input_dim = dims
    p.softmax.num_classes = vocab

    with self.session(use_gpu=False, graph=tf.Graph()) as sess:
      lm = p.Instantiate()
      np.random.seed(12345)
      inputs = np.random.randint(vocab, size=[time, batch])
      targets = np.zeros([time, batch])
      targets[:-1] = inputs[1:]
      inputs = tf.constant(inputs, tf.int32)
      paddings = np.zeros([time, batch])
      paddings[-1] = 1.0
      paddings = tf.constant(paddings, tf.float64)
      targets = tf.constant(targets, tf.int32)
      xent_output, _ = lm.FPropDefaultTheta(
          inputs=inputs,
          paddings=paddings,
          labels=py_utils.NestedMap(
              class_weights=1 - paddings, class_ids=targets))

      lm_vars = lm.vars.Flatten()
      grads = tf.gradients(xent_output.avg_xent, lm_vars)

      for i, x in enumerate(grads):
        if isinstance(x, tf.IndexedSlices):
          grads[i] = tf.math.unsorted_segment_sum(x.values, x.indices,
                                                  x.dense_shape[0])

      self.evaluate(tf.global_variables_initializer())
      self.assertEqual(len(lm_vars), len(grads))
      for x, grad_x in zip(lm_vars, grads):
        grad_symbolic = self.evaluate(grad_x)
        grad_numeric = test_utils.ComputeNumericGradient(
            sess, xent_output.avg_xent, x, delta=1e-6)
        self.assertAllClose(grad_symbolic, grad_numeric, atol=0.005)


class GPipeTransformerLmTest(test_utils.TestCase):

  def _testParams(self, batch, dims, hidden_dim, vocab):
    p = lm_layers.GPipeTransformerLm.Params()
    p.name = 'transformerlm'
    p.vocab_size = vocab
    p.stack.num_encoder_layers = 2
    p.stack.packed_input = False
    p.stack.emb_tpl.token_emb.vocab_size = vocab
    p.stack.emb_tpl.token_emb.embedding_dim = dims
    p.stack.emb_tpl.position_emb.embedding_dim = dims
    p.stack.model_dim = dims
    p.stack.splits = 2
    p.stack.softmax_tpl.input_dim = dims
    p.stack.softmax_tpl.num_classes = vocab
    p.stack.batch_dim = 0
    trans_tpl = batch_major_attention.GPipeBatchMajorTransformerLayer.Params()
    trans_tpl.has_aux_atten = False
    trans_tpl.mask_self_atten = True
    trans_tpl.input_dim = dims
    trans_tpl.output_dim = dims
    trans_tpl.tr_atten_tpl.input_dim = dims
    trans_tpl.tr_atten_tpl.hidden_dim = dims
    trans_tpl.tr_atten_tpl.is_masked = True
    trans_tpl.tr_atten_tpl.num_heads = 2
    trans_tpl.tr_fflayer_tpl.hidden_dim = hidden_dim
    p.stack.encoder_tpl = trans_tpl
    p.random_seed = 12345
    return p

  def _SetupGraph(self, p, time, batch, vocab, return_grad=False):
    lm = p.Instantiate()
    np.random.seed(12345)
    inputs = np.random.randint(vocab, size=[time, batch])
    targets = np.zeros([time, batch])
    targets[:-1] = inputs[1:]
    inputs = tf.constant(inputs, tf.int32)
    paddings = np.zeros([time, batch])
    paddings[-1] = 1.0
    paddings = tf.constant(paddings, tf.float64 if return_grad else tf.float32)
    targets = tf.constant(targets, tf.int32)
    xent_output, _ = lm.FPropDefaultTheta(
        inputs=inputs,
        paddings=paddings,
        labels=py_utils.NestedMap(
            class_weights=1 - paddings, class_ids=targets))
    if not return_grad:
      return xent_output

    lm_vars = lm.vars.Flatten()
    grads = tf.gradients(xent_output.avg_xent, lm_vars)
    for i, x in enumerate(grads):
      if isinstance(x, tf.IndexedSlices):
        grads[i] = tf.math.unsorted_segment_sum(x.values, x.indices,
                                                x.dense_shape[0])
    self.assertEqual(len(lm_vars), len(grads))
    return xent_output, lm_vars, grads

  def testBasic(self):
    time, batch, dims, hidden_dim, vocab = 5, 3, 6, 4, 8
    p = self._testParams(batch, dims, hidden_dim, vocab)
    xent_output = self._SetupGraph(p, time, batch, vocab)
    assert p.stack.encoder_tpl.tr_atten_tpl.is_masked
    with self.session():
      self.evaluate(tf.global_variables_initializer())
      xent_output_val = self.evaluate(xent_output)

      print('xent_output_val', xent_output_val)
      test_utils.CompareToGoldenSingleFloat(self, 3.083209,
                                            xent_output_val.avg_xent)
      self.assertAllEqual(xent_output_val.per_example_argmax,
                          np.argmax(xent_output_val.logits, axis=-1))

  def testBasicGrad(self):
    time, batch, dims, hidden_dim, vocab = 5, 3, 6, 4, 8
    p = self._testParams(batch, dims, hidden_dim, vocab)
    p.dtype = tf.float64
    xent_output, lm_vars, grads = self._SetupGraph(
        p, time, batch, vocab, return_grad=True)
    with self.session() as sess:
      self.evaluate(tf.global_variables_initializer())
      grads_symbolic = self.evaluate(grads)
      for x, grad_symbolic in zip(lm_vars, grads_symbolic):
        grad_numeric = test_utils.ComputeNumericGradient(
            sess, xent_output.avg_xent, x, delta=1e-6)
        self.assertAllClose(grad_symbolic, grad_numeric, atol=0.05)


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