# Copyright 2021 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 lingvo Jax embedding and softmax layers."""

from absl.testing import absltest
from absl.testing import parameterized
import jax
from jax import numpy as jnp
from lingvo.core import layers as lingvo_layers
from lingvo.jax import py_utils
from lingvo.jax import test_utils
from lingvo.jax.layers import embedding_softmax
import numpy as np
import tensorflow.compat.v2 as tf

to_np = test_utils.to_np


class EmbeddingSoftmaxTest(test_utils.TestCase):

  def setUp(self):
    super().setUp()
    np.random.seed(123456)
    tf.random.set_seed(123)

  @parameterized.parameters(('index', True), ('index', False), ('matmul', True),
                            ('matmul', False))
  def test_single_sharded_embedding_layer(self, lookup_style, scale_sqrt_depth):
    p = embedding_softmax.SingleShardEmbedding.Params().Set(
        name='jax_emb_lookup',
        vocab_size=10,
        embedding_dims=40,
        lookup_style=lookup_style,
        scale_sqrt_depth=scale_sqrt_depth)
    emb_layer = p.Instantiate()
    prng_key = jax.random.PRNGKey(seed=123)
    initial_vars = emb_layer.instantiate_variables(prng_key)
    npy_input = np.random.randint(0, p.vocab_size, [10, 20]).astype('int32')
    inputs = jnp.asarray(npy_input)
    outputs = test_utils.apply(emb_layer, initial_vars, emb_layer.fprop, inputs)
    # Test whether tf Embedding layer returns same output
    # Modify initial_vars to use TF compatible params
    tf_initial_vars = initial_vars
    tf_p = lingvo_layers.SingleShardEmbeddingLayer.Params().Set(
        name='tf_emb_lookup',
        vocab_size=p.vocab_size,
        embedding_dim=p.embedding_dims,
        scale_sqrt_depth=scale_sqrt_depth)
    tf_emb_layer = tf_p.Instantiate()
    tf_output = tf_emb_layer.FProp(tf_initial_vars,
                                   tf.constant(inputs, dtype=tf.int32))
    np_outputs = to_np(outputs)
    tf_np_outputs = to_np(tf_output)
    self.assertAllClose(tf_np_outputs, np_outputs, atol=1e-6)

  @parameterized.parameters((0., True, False, 0), (0., False, True, 0),
                            (1.0, True, False, 0.1), (1.0, False, True, 0.1))
  def test_single_sharded_softmax_layer(self, soft_cap_logits, use_class_ids,
                                        use_class_probabilities,
                                        label_smoothing_prob):
    if use_class_ids:
      class_ids = np.random.randint(0, 50, [8, 10, 1])
    else:
      class_ids = None
    if use_class_probabilities:
      class_probabilities = np.random.normal(1.5, 2.0, [8, 10, 50])
    else:
      class_probabilities = None
    p = embedding_softmax.SingleShardFullSoftmax.Params().Set(
        name='jax_softmax',
        num_classes=50,
        input_dims=40,
        soft_cap_logits=soft_cap_logits,
        label_smoothing_prob=label_smoothing_prob)
    softmax_layer = p.Instantiate()
    prng_key = jax.random.PRNGKey(seed=1234)
    initial_vars = softmax_layer.instantiate_variables(prng_key)
    npy_input = np.random.normal(1.5, 2.0, [8, 10, p.input_dims])
    inputs = jnp.asarray(npy_input)
    class_weights = np.random.normal(1.5, 2.0, [8, 10, 1])
    if class_probabilities is not None:
      class_probabilities /= np.sum(class_probabilities, axis=-1, keepdims=True)
    logits = test_utils.apply(softmax_layer, initial_vars,
                              softmax_layer.get_logits, inputs)
    outputs = test_utils.apply(
        softmax_layer,
        initial_vars,
        softmax_layer.fprop,
        inputs,
        class_weights,
        class_ids=class_ids,
        class_probabilities=class_probabilities)
    # Test whether tf Softmax layer returns same output
    # Modify initial_vars to use TF compatible params
    tf_initial_vars = initial_vars
    tf_initial_vars.linear = py_utils.NestedMap()
    tf_initial_vars.linear.w = initial_vars.logits_ffn.linear.w
    tf_initial_vars.bias = py_utils.NestedMap()
    tf_initial_vars.bias.b = initial_vars.logits_ffn.bias.b
    tf_p = lingvo_layers.SingleShardFullSoftmax.Params().Set(
        name='tf_softmax',
        num_classes=p.num_classes,
        input_dim=p.input_dims,
        logits_soft_max=soft_cap_logits)
    tf_softmax_layer = tf_p.Instantiate()
    tf_logits = tf_softmax_layer.Logits(tf_initial_vars,
                                        tf.constant(inputs, dtype=tf.float32))
    if use_class_ids and label_smoothing_prob > 0:
      class_probabilities = np.zeros([8, 10, 50])
      index = np.indices([8, 10])
      class_probabilities[index[0], index[1], np.squeeze(class_ids, 2)] = 1
      class_probabilities = (
          class_probabilities * (1 - label_smoothing_prob) +
          (1 - class_probabilities) * label_smoothing_prob /
          (p.num_classes - 1))
      class_ids = None
    tf_output = tf_softmax_layer.FProp(
        tf_initial_vars,
        tf.constant(inputs, dtype=tf.float32),
        class_weights,
        class_ids=class_ids,
        class_probabilities=class_probabilities)
    # Check all entries in the NestedMap and ensure it matches TF
    np_get_logits = to_np(logits)
    tf_np_get_logits = to_np(tf_logits)
    self.assertAllClose(np_get_logits, tf_np_get_logits, atol=1e-6)
    # Note: The argmax-related values are very sensitive to numerical errors.
    for k in outputs.keys():
      self.assertAllClose(to_np(outputs[k]), to_np(tf_output[k]), atol=1e-6)

  def test_simple_softmax_layer_class_ids(self):
    batch_size = 8
    num_classes = 50
    class_ids = np.random.randint(0, 50, [8, 1])
    p = embedding_softmax.SingleShardFullSoftmax.Params().Set(
        name='jax_softmax', num_classes=num_classes, input_dims=40)
    softmax_layer = p.Instantiate()
    prng_key = jax.random.PRNGKey(seed=123)
    initial_vars = softmax_layer.instantiate_variables(prng_key)
    npy_input = np.random.normal(1.5, 2.0, [batch_size, p.input_dims])
    inputs = jnp.asarray(npy_input)
    class_weights = np.random.normal(1.5, 2.0, [batch_size, 1])
    logits = test_utils.apply(softmax_layer, initial_vars,
                              softmax_layer.get_logits, inputs)
    outputs = test_utils.apply(
        softmax_layer,
        initial_vars,
        softmax_layer.fprop,
        inputs,
        class_weights,
        class_ids=class_ids,
        class_probabilities=None)
    # Test whether tf Softmax layer returns same output.
    # Modify initial_vars to use TF compatible params.
    tf_initial_vars = test_utils.replace_jax_simple_full_softmax_vars_to_tf(
        initial_vars)
    # Convert all the values to TF tensor.
    tf_initial_vars = tf.nest.map_structure(tf.convert_to_tensor,
                                            tf_initial_vars)

    tf_p = lingvo_layers.SimpleFullSoftmax.Params().Set(
        name='tf_softmax', num_classes=p.num_classes, input_dim=p.input_dims)
    tf_softmax_layer = tf_p.Instantiate()
    tf_logits = tf_softmax_layer.Logits(tf_initial_vars,
                                        tf.constant(inputs, dtype=tf.float32))
    tf_output = tf_softmax_layer.FProp(
        tf_initial_vars,
        tf.constant(inputs, dtype=tf.float32),
        class_weights,
        class_ids=class_ids,
        class_probabilities=None)
    # Check all entries in the NestedMap and ensure it matches TF.
    np_get_logits = to_np(logits)
    tf_np_get_logits = to_np(tf_logits)
    self.assertAllClose(np_get_logits, tf_np_get_logits)
    for k in outputs.keys():
      self.assertAllClose(to_np(outputs[k]), to_np(tf_output[k]))

  @parameterized.parameters((8, 1001), (16, 1024), (32, 30000))
  def test_simple_softmax_layer_class_probs(self, batch_size, num_classes):
    batch_size = 8
    num_classes = 1001
    class_probabilities = np.random.normal(1.5, 2.0, [batch_size, num_classes])
    # Normalize class probabilities to be a probability distribution.
    class_probabilities /= np.sum(class_probabilities, axis=-1, keepdims=True)
    p = embedding_softmax.SingleShardFullSoftmax.Params().Set(
        name='jax_softmax', num_classes=num_classes, input_dims=40)
    softmax_layer = p.Instantiate()
    prng_key = jax.random.PRNGKey(seed=123)
    initial_vars = softmax_layer.instantiate_variables(prng_key)
    npy_input = np.random.normal(1.5, 2.0, [batch_size, p.input_dims])
    inputs = jnp.asarray(npy_input)
    class_weights = np.random.normal(1.5, 2.0, [batch_size, 1])
    logits = test_utils.apply(softmax_layer, initial_vars,
                              softmax_layer.get_logits, inputs)
    outputs = test_utils.apply(
        softmax_layer,
        initial_vars,
        softmax_layer.fprop,
        inputs,
        class_weights,
        class_ids=None,
        class_probabilities=class_probabilities)
    # Test whether tf Softmax layer returns same output.
    # Modify initial_vars to use TF compatible params.
    tf_initial_vars = test_utils.replace_jax_simple_full_softmax_vars_to_tf(
        initial_vars)
    # Convert all the values to TF tensor.
    tf_initial_vars = tf.nest.map_structure(tf.convert_to_tensor,
                                            tf_initial_vars)

    tf_p = lingvo_layers.SimpleFullSoftmax.Params().Set(
        name='tf_softmax', num_classes=p.num_classes, input_dim=p.input_dims)
    tf_softmax_layer = tf_p.Instantiate()
    tf_logits = tf_softmax_layer.Logits(tf_initial_vars,
                                        tf.constant(inputs, dtype=tf.float32))
    tf_output = tf_softmax_layer.FProp(
        tf_initial_vars,
        tf.constant(inputs, dtype=tf.float32),
        class_weights,
        class_ids=None,
        class_probabilities=class_probabilities)
    # Check all entries in the NestedMap and ensure it matches TF.
    np_get_logits = to_np(logits)
    tf_np_get_logits = to_np(tf_logits)
    self.assertAllClose(np_get_logits, tf_np_get_logits)
    for k in outputs.keys():
      self.assertAllClose(to_np(outputs[k]), to_np(tf_output[k]))

  def test_simple_softmax_layer_value_error(self):
    batch_size = 8
    num_classes = 50
    class_ids = None
    class_probabilities = None
    p = embedding_softmax.SingleShardFullSoftmax.Params().Set(
        name='jax_softmax', num_classes=num_classes, input_dims=40)
    softmax_layer = p.Instantiate()
    prng_key = jax.random.PRNGKey(seed=123)
    initial_vars = softmax_layer.instantiate_variables(prng_key)
    npy_input = np.random.normal(1.5, 2.0, [batch_size, p.input_dims])
    inputs = jnp.asarray(npy_input)
    class_weights = np.random.normal(1.5, 2.0, [batch_size, 1])
    with self.assertRaises(ValueError):
      _ = test_utils.apply(
          softmax_layer,
          initial_vars,
          softmax_layer.fprop,
          inputs,
          class_weights,
          class_ids=class_ids,
          class_probabilities=class_probabilities)

  @parameterized.parameters((0., 'index', True), (0., 'matmul', True),
                            (1.0, 'index', False), (1.0, 'matmul', False))
  def test_single_sharded_shared_embedding_softmax_layer(
      self, soft_cap_logits, lookup_style, scale_sqrt_depth):
    class_ids = np.random.randint(1, 50, [8, 10, 1])
    p = embedding_softmax.SingleShardSharedEmbeddingSoftmax.Params().Set(
        name='jax_softmax',
        num_classes=50,
        input_dims=40,
        soft_cap_logits=soft_cap_logits,
        lookup_style=lookup_style,
        scale_sqrt_depth=scale_sqrt_depth)
    softmax_layer = p.Instantiate()
    prng_key = jax.random.PRNGKey(seed=123)
    initial_vars = softmax_layer.instantiate_variables(prng_key)
    npy_input = np.random.normal(1.5, 2.0, [8, 10, p.input_dims])
    inputs = jnp.asarray(npy_input)
    class_weights = np.random.normal(1.5, 2.0, [8, 10, 1])
    outputs = test_utils.apply(
        softmax_layer,
        initial_vars,
        softmax_layer.fprop,
        inputs,
        class_weights,
        class_ids=class_ids)
    ids = np.squeeze(class_ids, axis=-1)
    emb_lookup_outputs = test_utils.apply(
        softmax_layer,
        initial_vars,
        softmax_layer.emb_lookup,
        ids=jnp.asarray(ids))
    # Test whether tf Softmax layer returns same output
    # Modify initial_vars to use TF compatible params
    tf_initial_vars = initial_vars
    tf_initial_vars.linear = py_utils.NestedMap()
    tf_initial_vars.linear.w = initial_vars.logits_ffn.linear.w
    tf_initial_vars.bias = py_utils.NestedMap()
    tf_initial_vars.bias.b = initial_vars.logits_ffn.bias.b
    tf_p = lingvo_layers.SingleShardSharedEmbeddingSoftmax.Params().Set(
        name='tf_softmax',
        num_classes=p.num_classes,
        input_dim=p.input_dims,
        vocab_size=p.num_classes,
        embedding_dim=p.input_dims,
        logits_soft_max=soft_cap_logits,
        scale_sqrt_depth=scale_sqrt_depth)
    tf_softmax_layer = tf_p.Instantiate()
    tf_output = tf_softmax_layer.FProp(
        tf_initial_vars,
        tf.constant(inputs, dtype=tf.float32),
        class_weights,
        class_ids=class_ids)
    tf_emb_lookup_output = tf_softmax_layer.EmbLookup(
        tf_initial_vars, ids=tf.constant(ids))

    # Check all entries in the NestedMap and ensure it matches TF
    np_logits = to_np(outputs.logits)
    tf_np_logits = to_np(tf_output.logits)
    self.assertAllClose(np_logits, tf_np_logits, atol=1e-6)
    for k in outputs.keys():
      self.assertAllClose(to_np(outputs[k]), to_np(tf_output[k]), atol=1e-6)
    np_emb_lookup_output = to_np(emb_lookup_outputs)
    tf_np_emb_lookup_output = to_np(tf_emb_lookup_output)
    self.assertAllClose(
        tf_np_emb_lookup_output, np_emb_lookup_output, atol=1e-6)

  @parameterized.parameters((1, 10), (1, 1e5), (10, 20), (10, 1e5))
  def test_position_embedding_layer(self, min_timescale, max_timescale):
    p = embedding_softmax.PositionalEmbedding.Params().Set(
        name='jax_pos',
        embedding_dims=50,
        min_timescale=min_timescale,
        max_timescale=max_timescale)
    pos_layer = p.Instantiate()
    prng_key = jax.random.PRNGKey(seed=123)
    initial_vars = pos_layer.instantiate_variables(prng_key)
    seq_length = np.random.randint(100, 1000)
    output = test_utils.apply(pos_layer, initial_vars, pos_layer.fprop,
                              seq_length)
    output = jnp.squeeze(output, axis=0)
    # Test whether tf PositionalEmbedding layer returns same output
    # Modify initial_vars to use TF compatible params
    tf_initial_vars = initial_vars
    tf_p = lingvo_layers.PositionalEmbeddingLayer.Params().Set(
        name='tf_pos',
        embedding_dim=p.embedding_dims,
        min_timescale=min_timescale,
        max_timescale=max_timescale)
    tf_pos_layer = tf_p.Instantiate()
    tf_output = tf_pos_layer.FProp(tf_initial_vars, seq_length)
    np_pos = to_np(output)
    tf_np_pos = to_np(tf_output)
    self.assertAllClose(tf_np_pos, np_pos, atol=1e-3)

  @parameterized.parameters((1, 10), (1, 1e5), (10, 20), (10, 1e5))
  def test_position_embedding_layer_with_position(self, min_timescale,
                                                  max_timescale):
    p = embedding_softmax.PositionalEmbedding.Params().Set(
        name='jax_pos',
        embedding_dims=50,
        min_timescale=min_timescale,
        max_timescale=max_timescale)
    pos_layer = p.Instantiate()
    prng_key = jax.random.PRNGKey(seed=123)
    initial_vars = pos_layer.instantiate_variables(prng_key)
    position = np.array([[0, 1, 2, 3, 4, 0, 1, 2, 3, 4],
                         [0, 1, 2, 0, 1, 2, 0, 1, 2, 0],
                         [0, 1, 2, 3, 4, 5, 6, 0, 1, 2],
                         [0, 1, 2, 3, 4, 5, 6, 7, 8, 9]])
    output = test_utils.apply(
        pos_layer, initial_vars, pos_layer.fprop, position=position)
    # Test whether tf PositionalEmbedding layer returns same output
    # Modify initial_vars to use TF compatible params
    tf_initial_vars = initial_vars
    tf_p = lingvo_layers.PositionalEmbeddingLayer.Params().Set(
        name='tf_pos',
        embedding_dim=p.embedding_dims,
        min_timescale=min_timescale,
        max_timescale=max_timescale)
    tf_pos_layer = tf_p.Instantiate()
    tf_output = tf_pos_layer.FPropWithPosition(tf_initial_vars, position)
    np_pos = to_np(output)
    tf_np_pos = to_np(tf_output)
    self.assertAllClose(tf_np_pos, np_pos, atol=1e-3)

  @parameterized.parameters((1, 10, 1), (1, 1e5, 3), (10, 20, 4), (10, 1e5, 5))
  def test_rotary_position_embedding_layer_prefix(self, min_timescale,
                                                  max_timescale, window_size):
    embedding_dims = 32
    p = embedding_softmax.RotaryPositionalEmbedding.Params().Set(
        name='jax_pos',
        embedding_dims=embedding_dims,
        min_timescale=min_timescale,
        max_timescale=max_timescale)
    pos_layer = p.Instantiate()
    prng_key = jax.random.PRNGKey(seed=123)
    initial_vars = pos_layer.instantiate_variables(prng_key)
    inputs = np.random.normal(1.5, 2.5, (2, 8, 4, embedding_dims))
    output = test_utils.apply(
        pos_layer, initial_vars, pos_layer.fprop, inputs=inputs)
    # Test whether extend_step returns same output.
    for i in range(inputs.shape[1]):
      start = max(0, i + 1 - window_size)
      end = i + 1
      inputs_prefix = inputs[:, start:end, :, :]
      pad_width = window_size - end + start
      paddings = [(0, 0), (pad_width, 0), (0, 0), (0, 0)]
      inputs_prefix = jnp.pad(inputs_prefix, paddings)
      jax_extend_step_out = test_utils.apply(
          pos_layer,
          initial_vars,
          pos_layer.extend_step,
          inputs_prefix,
          time_step=i)
      jax_extend_step_out = jax.lax.dynamic_slice_in_dim(
          jax_extend_step_out,
          start_index=window_size - 1,
          slice_size=1,
          axis=1)
      jax_np_extend_step_out = test_utils.to_np(jax_extend_step_out)
      jax_fprop_slice = jax.lax.dynamic_slice_in_dim(
          output, start_index=i, slice_size=1, axis=1)
      self.assertArraysEqual(jax_fprop_slice, jax_np_extend_step_out)

  @parameterized.parameters((1, 10), (1, 1e5), (10, 20), (10, 1e5))
  def test_rotary_position_embedding_layer_no_prefix(self, min_timescale,
                                                     max_timescale):
    embedding_dims = 32
    p = embedding_softmax.RotaryPositionalEmbedding.Params().Set(
        name='jax_pos',
        embedding_dims=embedding_dims,
        min_timescale=min_timescale,
        max_timescale=max_timescale)
    pos_layer = p.Instantiate()
    prng_key = jax.random.PRNGKey(seed=123)
    initial_vars = pos_layer.instantiate_variables(prng_key)
    inputs = np.random.normal(1.5, 2.5, (2, 8, 4, embedding_dims))
    output = test_utils.apply(
        pos_layer, initial_vars, pos_layer.fprop, inputs=inputs)
    # Test whether extend_step returns same output.
    for i in range(inputs.shape[1]):
      jax_extend_step_out = test_utils.apply(
          pos_layer,
          initial_vars,
          pos_layer.extend_step,
          inputs[:, i, :, :],
          time_step=i)
      jax_np_extend_step_out = test_utils.to_np(jax_extend_step_out)
      jax_fprop_slice = output[:, i, :, :]
      self.assertArraysEqual(jax_fprop_slice, jax_np_extend_step_out)

  @parameterized.parameters(
      ([0, 1, 0, 1],),
      ([0, 1, 2, 3],),
      ([0, 1, 2, 0],),
      ([0, 0, 1, 2],),
      (None),
  )
  def test_rotary_position_embedding_layer_2d(self, position):
    embedding_dims = 2
    min_timescale = 1
    max_timescale = 1e4
    p = embedding_softmax.RotaryPositionalEmbedding.Params().Set(
        name='jax_pos',
        embedding_dims=embedding_dims,
        min_timescale=min_timescale,
        max_timescale=max_timescale)
    pos_layer = p.Instantiate()
    prng_key = jax.random.PRNGKey(seed=123)
    initial_vars = pos_layer.instantiate_variables(prng_key)
    inputs = np.random.normal(1.5, 2.5, (1, 4, 1, embedding_dims))
    if position is None:
      position = jnp.arange(4, dtype=jnp.float32)
    position = jnp.array(position)
    output = test_utils.apply(
        pos_layer,
        initial_vars,
        pos_layer.fprop,
        inputs=inputs,
        position=position[jnp.newaxis, :])
    np_output = test_utils.to_np(output)
    sinusoid_inp = position
    sin = jnp.sin(sinusoid_inp)
    cos = jnp.cos(sinusoid_inp)
    first_part = inputs[0, :, 0, 0] * cos - inputs[0, :, 0, 1] * sin
    second_part = inputs[0, :, 0, 1] * cos + inputs[0, :, 0, 0] * sin
    expected_output = np.stack([first_part, second_part], axis=-1)
    self.assertArraysEqual(np_output[0, :, 0, :], expected_output)

  @parameterized.parameters('index', 'matmul')
  def test_trainable_positional_embedding_layer(self, lookup_style):
    p = embedding_softmax.TrainablePositionalEmbedding.Params().Set(
        name='jax_pos_emb',
        max_seq_length=10,
        embedding_dims=40,
        lookup_style=lookup_style)
    emb_layer = p.Instantiate()
    prng_key = jax.random.PRNGKey(seed=123)
    initial_vars = emb_layer.instantiate_variables(prng_key)
    npy_input = np.random.randint(0, p.max_seq_length,
                                  [10, p.max_seq_length]).astype('int32')
    inputs = jnp.asarray(npy_input)
    outputs = test_utils.apply(emb_layer, initial_vars, emb_layer.fprop,
                               p.max_seq_length, inputs)
    # Test whether tf Embedding layer returns same output
    # Modify initial_vars to use TF compatible params
    tf_initial_vars = initial_vars
    tf_p = lingvo_layers.SingleShardEmbeddingLayer.Params().Set(
        name='tf_pos_emb',
        vocab_size=p.max_seq_length,
        embedding_dim=p.embedding_dims)
    tf_emb_layer = tf_p.Instantiate()
    tf_output = tf_emb_layer.FProp(tf_initial_vars,
                                   tf.constant(inputs, dtype=tf.int32))
    np_outputs = to_np(outputs)
    tf_np_outputs = to_np(tf_output)
    self.assertAllClose(tf_np_outputs, np_outputs)

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