# Copyright 2019 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 dense attention layers."""

import tensorflow.compat.v2 as tf

from absl.testing import parameterized
import numpy as np

import keras
from keras import combinations
from keras import testing_utils
from keras.layers import core
from keras.layers import dense_attention
from keras.mixed_precision import policy


@combinations.generate(combinations.combine(mode=['graph', 'eager']))
class BaseDenseAttentionTest(tf.test.TestCase, parameterized.TestCase):

  def test_one_dim_with_mask(self):
    # Scores tensor of shape [1, 1, 1]
    scores = np.array([[[1.1]]], dtype=np.float32)
    # Value tensor of shape [1, 1, 1]
    v = np.array([[[1.6]]], dtype=np.float32)
    # Scores mask tensor of shape [1, 1, 1]
    scores_mask = np.array([[[True]]], dtype=np.bool_)
    actual, actual_scores = dense_attention.BaseDenseAttention()._apply_scores(
        scores=scores, value=v, scores_mask=scores_mask)

    # Expected softmax_scores = [[[1]]]
    expected_scores = np.array([[[1.]]], dtype=np.float32)
    self.assertAllClose(expected_scores, actual_scores)
    # Expected tensor of shape [1, 1, 1].
    # expected000 = softmax_scores[0, 0] * 1.6 = 1.6
    expected = np.array([[[1.6]]], dtype=np.float32)
    self.assertAllClose(expected, actual)

  def test_one_dim_no_mask(self):
    # Scores tensor of shape [1, 1, 1]
    scores = np.array([[[1.1]]], dtype=np.float32)
    # Value tensor of shape [1, 1, 1]
    v = np.array([[[1.6]]], dtype=np.float32)
    actual, actual_scores = dense_attention.BaseDenseAttention()._apply_scores(
        scores=scores, value=v)

    # Expected softmax_scores = [[[1]]]
    expected_scores = np.array([[[1.]]], dtype=np.float32)
    self.assertAllClose(expected_scores, actual_scores)
    # Expected tensor of shape [1, 1, 1].
    # expected000 = softmax_scores[0, 0] * 1.6 = 1.6
    expected = np.array([[[1.6]]], dtype=np.float32)
    self.assertAllClose(expected, actual)

  def test_multi_dim_with_mask(self):
    # Scores tensor of shape [1, 1, 3]
    scores = np.array([[[1., 0., 1.]]], dtype=np.float32)
    # Value tensor of shape [1, 3, 1]
    v = np.array([[[1.6], [0.7], [-0.8]]], dtype=np.float32)
    # Scores mask tensor of shape [1, 1, 3]
    scores_mask = np.array([[[True, True, False]]], dtype=np.bool_)
    actual, actual_scores = dense_attention.BaseDenseAttention()._apply_scores(
        scores=scores, value=v, scores_mask=scores_mask)

    # Expected softmax scores = softmax(scores) with zeros in positions where
    # v_mask == False.
    # => softmax_scores000 = exp(1)/(exp(1) + exp(0)) = 0.73105857863
    #    softmax_scores001 = exp(0)/(exp(1) + exp(0)) = 0.26894142137
    #    softmax_scores002 = 0
    expected_scores = np.array([[[0.73105857863, 0.26894142137, 0.]]],
                               dtype=np.float32)
    self.assertAllClose(expected_scores, actual_scores)
    # Expected tensor of shape [1, 1, 1].
    # expected000 = 0.73105857863 * 1.6 + 0.26894142137 * 0.7 - 0 * 0.8
    #             = 1.35795272077
    expected = np.array([[[1.35795272077]]], dtype=np.float32)
    self.assertAllClose(expected, actual)

  def test_multi_dim_no_mask(self):
    # Scores tensor of shape [1, 1, 3]
    scores = np.array([[[1., 0., 1.]]], dtype=np.float32)
    # Value tensor of shape [1, 3, 1]
    v = np.array([[[1.6], [0.7], [-0.8]]], dtype=np.float32)
    actual, actual_scores = dense_attention.BaseDenseAttention()._apply_scores(
        scores=scores, value=v)

    # Expected softmax_scores = softmax(scores).
    # => softmax_scores000 = exp(1)/(exp(1) + exp(0) + exp(1))
    #                      = 0.42231879825
    #    softmax_scores001 = exp(0)/(exp(1) + exp(0) + exp(1))
    #                      = 0.15536240349
    #    softmax_scores002 = exp(1)/(exp(1) + exp(0) + exp(1))
    #                      = 0.42231879825
    expected_scores = np.array(
        [[[0.42231879825, 0.15536240349, 0.42231879825]]], dtype=np.float32)
    self.assertAllClose(expected_scores, actual_scores)
    # Expected tensor of shape [1, 1, 1].
    # expected000 = 0.42231879825 * 1.6 + 0.15536240349 * 0.7
    #               - 0.42231879825 * 0.8
    #             = 0.44660872104
    expected = np.array([[[0.44660872104]]], dtype=np.float32)
    self.assertAllClose(expected, actual)

  def test_one_dim_batch_size_two(self):
    # Scores tensor of shape [2, 1, 1]
    scores = np.array([[[1.1]], [[2.1]]], dtype=np.float32)
    # Value tensor of shape [2, 1, 1]
    v = np.array([[[1.6]], [[2.6]]], dtype=np.float32)
    # Scpres mask tensor of shape [2, 1, 1]
    scores_mask = np.array([[[True]], [[True]]], dtype=np.bool_)
    actual, actual_scores = dense_attention.BaseDenseAttention()._apply_scores(
        scores=scores, value=v, scores_mask=scores_mask)

    # Expected softmax_scores = [[[1]], [[1]]]
    expected_scores = np.array([[[1.]], [[1.]]], dtype=np.float32)
    self.assertAllClose(expected_scores, actual_scores)
    # Expected tensor of shape [2, 1, 1].
    # expected000 = softmax_scores[0, 0] * 1.6 = 1.6
    # expected100 = softmax_scores[1, 0] * 2.6 = 2.6
    expected = np.array([[[1.6]], [[2.6]]], dtype=np.float32)
    self.assertAllClose(expected, actual)

  def test_shape_with_dropout(self):
    # scores: Scores float tensor of shape `[batch_size, tq, tv]`.
    # value: Value tensor of shape `[batch_size, tv, dim]`.
    batch_size = 4
    tq = 5
    tv = 6
    dim = 7
    scores = np.ones((batch_size, tq, tv))
    value = np.ones((batch_size, tv, dim))
    actual, actual_scores = dense_attention.BaseDenseAttention(
        dropout=0.1)._apply_scores(
            scores=scores, value=value, training=False)

    # Expected Tensor of shape `[batch_size, tq, tv]`.
    expected_scores_shape = [batch_size, tq, tv]
    self.assertAllEqual(expected_scores_shape, tf.shape(actual_scores))
    # Expected Tensor of shape `[batch_size, tq, dim]`.
    expected_shape = [batch_size, tq, dim]
    self.assertAllEqual(expected_shape, tf.shape(actual))

  def test_serialization(self):
    # Test serialization with causal
    layer = dense_attention.BaseDenseAttention(causal=True)

    config = keras.layers.serialize(layer)
    new_layer = keras.layers.deserialize(config)
    self.assertEqual(new_layer.causal, True)

    config = layer.get_config()
    new_layer = dense_attention.BaseDenseAttention.from_config(config)
    self.assertEqual(new_layer.causal, True)


@combinations.generate(combinations.combine(mode=['graph', 'eager']))
class AttentionTest(tf.test.TestCase, parameterized.TestCase):

  def test_calculate_scores_one_dim(self):
    # Query tensor of shape [1, 1, 1]
    q = np.array([[[1.1]]], dtype=np.float32)
    # Key tensor of shape [1, 1, 1]
    k = np.array([[[1.6]]], dtype=np.float32)
    attention_layer = dense_attention.Attention()
    attention_layer.build(input_shape=([1, 1, 1], [1, 1, 1]))
    actual = attention_layer._calculate_scores(query=q, key=k)

    # Expected tensor of shape [1, 1, 1].
    # expected000 = 1.1*1.6 = 1.76
    expected = np.array([[[1.76]]], dtype=np.float32)
    self.assertAllClose(expected, actual)

  def test_calculate_scores_multi_dim(self):
    # Query tensor of shape [1, 2, 4]
    q = np.array([[[1., 1.1, 1.2, 1.3], [2., 2.1, 2.2, 2.3]]], dtype=np.float32)
    # Key tensor of shape [1, 3, 4]
    k = np.array(
        [[[1.5, 1.6, 1.7, 1.8], [2.5, 2.6, 2.7, 2.8], [3.5, 3.6, 3.7, 3.8]]],
        dtype=np.float32)
    attention_layer = dense_attention.Attention()
    attention_layer.build(input_shape=([1, 2, 4], [1, 3, 4]))
    actual = attention_layer._calculate_scores(query=q, key=k)

    # Expected tensor of shape [1, 2, 3].
    # expected000 = 1.*1.5+1.1*1.6+1.2*1.7+1.3*1.8 = 7.64
    # expected001 = 1.*2.5+1.1*2.6+1.2*2.7+1.3*2.8 = 12.24
    # expected002 = 1.*3.5+1.1*3.6+1.2*3.7+1.3*3.8 = 16.84
    # expected010 = 2.*1.5+2.1*1.6+2.2*1.7+2.3*1.8 = 14.24
    # expected011 = 2.*2.5+2.1*2.6+2.2*2.7+2.3*2.8 = 22.84
    # expected012 = 2.*3.5+2.1*3.6+2.2*3.7+2.3*3.8 = 31.44
    expected = np.array([[[7.64, 12.24, 16.84], [14.24, 22.84, 31.44]]],
                        dtype=np.float32)
    self.assertAllClose(expected, actual)

  def test_calculate_scores_one_dim_batch_size_two(self):
    # Query tensor of shape [2, 1, 1]
    q = np.array([[[1.1]], [[2.1]]], dtype=np.float32)
    # Key tensor of shape [2, 1, 1]
    k = np.array([[[1.6]], [[2.6]]], dtype=np.float32)
    attention_layer = dense_attention.Attention()
    attention_layer.build(input_shape=([2, 1, 1], [2, 1, 1]))
    actual = attention_layer._calculate_scores(query=q, key=k)

    # Expected tensor of shape [2, 1, 1].
    # expected000 = 1.1*1.6 = 1.76
    # expected100 = 2.1*2.6 = 5.46
    expected = np.array([[[1.76]], [[5.46]]], dtype=np.float32)
    self.assertAllClose(expected, actual)

  def test_calculate_scores_one_dim_with_scale(self):
    """Tests that scores are multiplied by scale."""
    # Query tensor of shape [1, 1, 1]
    q = np.array([[[1.1]]], dtype=np.float32)
    # Key tensor of shape [1, 1, 1]
    k = np.array([[[1.6]]], dtype=np.float32)
    attention_layer = dense_attention.Attention(use_scale=True)
    attention_layer.build(input_shape=([1, 1, 1], [1, 1, 1]))
    attention_layer.scale = -2.
    actual = attention_layer._calculate_scores(query=q, key=k)

    # Expected tensor of shape [1, 1, 1].
    # expected000 = -2*1.1*1.6 = -3.52
    expected = np.array([[[-3.52]]], dtype=np.float32)
    self.assertAllClose(expected, actual)

  def test_shape(self):
    # Query tensor of shape [1, 2, 4]
    q = np.array([[[1., 1.1, 1.2, 1.3], [2., 2.1, 2.2, 2.3]]], dtype=np.float32)
    # Value tensor of shape [1, 3, 4]
    v = np.array(
        [[[1.5, 1.6, 1.7, 1.8], [2.5, 2.6, 2.7, 2.8], [3.5, 3.6, 3.7, 3.8]]],
        dtype=np.float32)
    # Value mask tensor of shape [1, 3]
    v_mask = np.array([[True, True, False]], dtype=np.bool_)
    attention_layer = dense_attention.Attention()
    actual = attention_layer([q, v], mask=[None, v_mask])

    expected_shape = [1, 2, 4]
    self.assertAllEqual(expected_shape, tf.shape(actual))

  def test_shape_with_key(self):
    # Query tensor of shape [1, 2, 4]
    q = np.array([[[1., 1.1, 1.2, 1.3], [2., 2.1, 2.2, 2.3]]], dtype=np.float32)
    # Value tensor of shape [1, 3, 4]
    v = np.array(
        [[[1.5, 1.6, 1.7, 1.8], [2.5, 2.6, 2.7, 2.8], [3.5, 3.6, 3.7, 3.8]]],
        dtype=np.float32)
    # Key tensor of shape [1, 3, 4]
    k = np.array(
        [[[1.5, 1.6, 1.7, 1.8], [2.5, 2.6, 2.7, 2.8], [3.5, 3.6, 3.7, 3.8]]],
        dtype=np.float32)
    # Value mask tensor of shape [1, 3]
    v_mask = np.array([[True, True, False]], dtype=np.bool_)
    attention_layer = dense_attention.Attention()
    actual = attention_layer([q, v, k], mask=[None, v_mask])

    expected_shape = [1, 2, 4]
    self.assertAllEqual(expected_shape, tf.shape(actual))

  def test_multi_dim(self):
    # Query tensor of shape [1, 1, 1]
    q = np.array([[[1.1]]], dtype=np.float32)
    # Value tensor of shape [1, 3, 1]
    v = np.array([[[1.6], [0.7], [-0.8]]], dtype=np.float32)
    # Value mask tensor of shape [1, 3]
    v_mask = np.array([[True, True, False]], dtype=np.bool_)
    attention_layer = dense_attention.Attention()
    actual = attention_layer([q, v], mask=[None, v_mask])

    # Expected scores of shape [1, 1, 3]
    # scores = [[[1.1*1.6, 1.1*0.7, -1.1*0.8]]] = [[[1.76, 0.77, -0.88]]]
    # Expected attention distribution = softmax(scores) with zeros in
    # positions where v_mask == False.
    # => attention_distribution000 = exp(1.76)/(exp(1.76) + exp(0.77))
    #                              = 0.72908792234
    #    attention_distribution001 = exp(0.77)/(exp(1.76) + exp(0.77))
    #                              = 0.27091207765
    #    attention_distribution002 = 0
    #
    # Expected tensor of shape [1, 1, 1].
    # expected000 = 0.72908792234 * 1.6 + 0.27091207765 * 0.7 - 0 * 0.8
    #             = 1.3561791301
    expected = np.array([[[1.3561791301]]], dtype=np.float32)
    self.assertAllClose(expected, actual)

  def test_multi_dim_with_key(self):
    # Query tensor of shape [1, 1, 1]
    q = np.array([[[1.1]]], dtype=np.float32)
    # Value tensor of shape [1, 3, 1]
    v = np.array([[[0.5], [0.8], [-0.3]]], dtype=np.float32)
    # Key tensor of shape [1, 3, 1]
    k = np.array([[[1.6], [0.7], [-0.8]]], dtype=np.float32)
    # Value mask tensor of shape [1, 3]
    v_mask = np.array([[True, True, False]], dtype=np.bool_)
    attention_layer = dense_attention.Attention()
    actual = attention_layer([q, v, k], mask=[None, v_mask])

    # Expected scores of shape [1, 1, 3]
    # scores = [[[1.1*1.6, 1.1*0.7, -1.1*0.8]]] = [[[1.76, 0.77, -0.88]]]
    # Expected attention distribution = softmax(scores) with zeros in
    # positions where v_mask == False.
    # => attention_distribution000 = exp(1.76)/(exp(1.76) + exp(0.77))
    #                              = 0.72908792234
    #    attention_distribution001 = exp(0.77)/(exp(1.76) + exp(0.77))
    #                              = 0.27091207765
    #    attention_distribution002 = 0
    #
    # Expected tensor of shape [1, 1, 1].
    # expected000 = 0.72908792234 * 0.5 + 0.27091207765 * 0.8 - 0 * 0.3
    #             = 0.58127362329
    expected = np.array([[[0.58127362329]]], dtype=np.float32)
    self.assertAllClose(expected, actual)

  @parameterized.named_parameters(
      ('', False),
      ('return_attention_scores', True),
  )
  def test_multi_dim_with_query_mask(self, return_attention_scores):
    # Query tensor of shape [1, 2, 1]
    q = np.array([[[1.1], [-0.5]]], dtype=np.float32)
    # Value tensor of shape [1, 3, 1]
    v = np.array([[[1.6], [0.7], [-0.8]]], dtype=np.float32)
    # Query mask tensor of shape [1, 2]
    q_mask = np.array([[True, False]], dtype=np.bool_)
    # Value mask tensor of shape [1, 3]
    v_mask = np.array([[True, True, False]], dtype=np.bool_)
    attention_layer = dense_attention.Attention()
    if return_attention_scores:
      actual, actual_scores = attention_layer(
          [q, v],
          mask=[q_mask, v_mask],
          return_attention_scores=return_attention_scores)
    else:
      actual = attention_layer([q, v],
                               mask=[q_mask, v_mask],
                               return_attention_scores=return_attention_scores)

    # Expected scores of shape [1, 2, 3]
    # scores = [[[1.1*1.6, 1.1*0.7, -1.1*0.8], [-0.5*1.6, -0.5*0.7, 0.5*0.8]]]
    #        = [[[1.76, 0.77, -0.88], [-0.8, -0.35, 0.4]]]
    # Expected attention distribution = softmax(scores) with zeros in
    # positions where v_mask == False.
    # => attention_distribution000 = exp(1.76)/(exp(1.76) + exp(0.77))
    #                              = 0.72908792234
    #    attention_distribution001 = exp(0.77)/(exp(1.76) + exp(0.77))
    #                              = 0.27091207765
    #    attention_distribution002 = 0
    # => attention_distribution010 = exp(-0.8)/(exp(-0.8) + exp(-0.35))
    #                              = 0.38936076605
    #    attention_distribution011 = exp(-0.35)/(exp(-0.8) + exp(-0.35))
    #                              = 0.61063923394
    #    attention_distribution012 = 0
    if return_attention_scores:
      expected_scores = np.array([[[0.72908792234, 0.27091207765, 0.],
                                   [0.38936076605, 0.61063923394, 0.]]],
                                 dtype=np.float32)
      self.assertAllClose(expected_scores, actual_scores)
    # Expected tensor of shape [1, 2, 1] with zeros where  q_mask == False.
    # expected000 = 0.72908792234 * 1.6 + 0.27091207765 * 0.7 - 0 * 0.8
    #             = 1.3561791301
    # expected000 = 0
    expected = np.array([[[1.3561791301], [0.]]], dtype=np.float32)
    self.assertAllClose(expected, actual)

  def test_scale_None(self):
    """Tests that scale is None by default."""
    attention_layer = dense_attention.Attention()
    attention_layer.build(input_shape=([1, 1, 1], [1, 1, 1]))
    self.assertIsNone(attention_layer.scale)

  def test_scale_init_eager(self):
    """Tests that scale initializes to 1 when use_scale=True."""
    if not tf.executing_eagerly():
      self.skipTest('Only run in eager mode')
    attention_layer = dense_attention.Attention(use_scale=True)
    attention_layer.build(input_shape=([1, 1, 1], [1, 1, 1]))
    self.assertAllClose(1., attention_layer.scale.value())

  def test_scale_init_graph(self):
    """Tests that scale initializes to 1 when use_scale=True."""
    with self.cached_session() as sess:
      attention_layer = dense_attention.Attention(use_scale=True)
      attention_layer.build(input_shape=([1, 1, 1], [1, 1, 1]))
      sess.run(attention_layer.scale.initializer)
      self.assertAllClose(1., attention_layer.scale.value())

  @parameterized.named_parameters(
      ('', False),
      ('return_attention_scores', True),
  )
  def test_self_attention_causal(self, return_attention_scores):
    # Query-value tensor of shape [1, 3, 1]
    q = np.array([[[0.5], [0.8], [-0.3]]], dtype=np.float32)
    attention_layer = dense_attention.Attention(causal=True)
    if return_attention_scores:
      actual, actual_scores = attention_layer(
          [q, q], return_attention_scores=return_attention_scores)
    else:
      actual = attention_layer([q, q],
                               return_attention_scores=return_attention_scores)

    # Expected scores of shape [1, 3, 3]
    # scores = [[0.25, 0.4, -0.15], [0.4, 0.64, -0.24], [-0.15, -0.24, 0.09]]
    # Expected attention distribution = softmax(scores) lower triangular
    # => attention_distribution00 = [1., 0., 0.]
    #    attention_distribution01
    #      = [exp(0.4), exp(0.64), 0.] / (exp(0.4) + exp(0.64))
    #      = [0.44028635073, 0.55971364926, 0.]
    #    attention_distribution02
    #      = [exp(-0.15), exp(-0.24), exp(0.09)]
    #        / (exp(-0.15) + exp(-0.24) + exp(0.09))
    #      = [0.31395396638, 0.28693232061, 0.399113713]
    if return_attention_scores:
      expected_scores = np.array(
          [[[1., 0., 0.], [0.44028635073, 0.55971364926, 0.],
            [0.31395396638, 0.28693232061, 0.399113713]]],
          dtype=np.float32)
      self.assertAllClose(expected_scores, actual_scores)
    # Expected tensor of shape [1, 3, 1].
    # expected000 = 0.5
    # expected010 = 0.44028635073 * 0.5 + 0.55971364926 * 0.8
    #             = 0.66791409477
    # expected020 = 0.31395396638 * 0.5 +0.28693232061 * 0.8 -0.399113713 * 0.3
    #             = 0.26678872577
    expected = np.array([[[0.5], [0.66791409477], [0.26678872577]]],
                        dtype=np.float32)
    self.assertAllClose(expected, actual)

  def test_inputs_not_list(self):
    attention_layer = dense_attention.Attention()
    q = np.array([[[1.1]]], dtype=np.float32)
    with self.assertRaisesRegex(
        ValueError, 'Attention layer must be called on a list of inputs'):
      attention_layer(q)

  def test_inputs_too_short(self):
    attention_layer = dense_attention.Attention()
    q = np.array([[[1.1]]], dtype=np.float32)
    with self.assertRaisesRegex(
        ValueError, 'Attention layer accepts inputs list of length 2 or 3'):
      attention_layer([q])

  def test_inputs_too_long(self):
    attention_layer = dense_attention.Attention()
    q = np.array([[[1.1]]], dtype=np.float32)
    with self.assertRaisesRegex(
        ValueError, 'Attention layer accepts inputs list of length 2 or 3'):
      attention_layer([q, q, q, q])

  def test_mask_not_list(self):
    attention_layer = dense_attention.Attention()
    q = np.array([[[1.1]]], dtype=np.float32)
    mask = np.array([[True]], dtype=np.bool_)
    with self.assertRaisesRegex(ValueError,
                                'Attention layer mask must be a list'):
      attention_layer([q, q], mask=mask)

  def test_mask_too_short(self):
    attention_layer = dense_attention.Attention()
    q = np.array([[[1.1]]], dtype=np.float32)
    mask = np.array([[True]], dtype=np.bool_)
    with self.assertRaisesRegex(
        ValueError, 'Attention layer mask must be a list of length 2'):
      attention_layer([q, q], mask=[mask])

  def test_mask_too_long(self):
    attention_layer = dense_attention.Attention()
    q = np.array([[[1.1]]], dtype=np.float32)
    mask = np.array([[True]], dtype=np.bool_)
    with self.assertRaisesRegex(
        ValueError, 'Attention layer mask must be a list of length 2'):
      attention_layer([q, q], mask=[mask, mask, mask])

  def test_override_mask(self):
    attention_layer = dense_attention.Attention()
    q = core.Masking()(np.array([[[1.1]]], dtype=np.float32))
    mask = np.array([[False]], dtype=np.bool_)
    actual = attention_layer([q, q], mask=[mask, mask])
    self.assertAllClose([[[0]]], actual)

  def test_implicit_mask(self):
    attention_layer = dense_attention.Attention()
    q = core.Masking(1.1)(np.array([[[1.1], [1]]], dtype=np.float32))
    v = core.Masking(1.2)(np.array([[[1.2], [1]]], dtype=np.float32))
    actual = attention_layer([q, v])
    self.assertAllClose([[[0], [1]]], actual)

  @parameterized.named_parameters(
      ('', False),
      ('use_scale', True),
  )
  def test_serialization(self, use_scale):
    # Test serialization with use_scale
    layer = dense_attention.Attention(use_scale=use_scale)

    config = keras.layers.serialize(layer)
    new_layer = keras.layers.deserialize(config)
    self.assertEqual(new_layer.use_scale, use_scale)

    config = layer.get_config()
    new_layer = dense_attention.Attention.from_config(config)
    self.assertEqual(new_layer.use_scale, use_scale)


@combinations.generate(combinations.combine(mode=['graph', 'eager']))
class AdditiveAttentionTest(tf.test.TestCase, parameterized.TestCase):

  def test_calculate_scores_one_dim(self):
    # Query tensor of shape [1, 1, 1]
    q = np.array([[[1.1]]], dtype=np.float32)
    # Key tensor of shape [1, 1, 1]
    k = np.array([[[1.6]]], dtype=np.float32)
    attention_layer = dense_attention.AdditiveAttention()
    attention_layer.build(input_shape=([1, 1, 1], [1, 1, 1]))
    # Scale tensor of shape [1]
    attention_layer.scale = np.array([[[0.5]]], dtype=np.float32)
    actual = attention_layer._calculate_scores(query=q, key=k)

    # Expected tensor of shape [1, 1, 1].
    # expected000 = 0.5 * tanh(1.1 + 1.6) = 0.49550372683
    expected = np.array([[[0.49550372683]]], dtype=np.float32)
    self.assertAllClose(expected, actual)

  def test_calculate_scores_multi_dim(self):
    # Query tensor of shape [1, 2, 4]
    q = np.array([[[1., 1.1, 1.2, 1.3], [2., 2.1, 2.2, 2.3]]], dtype=np.float32)
    # Key tensor of shape [1, 3, 4]
    k = np.array(
        [[[1.5, 1.6, 1.7, 1.8], [2.5, 2.6, 2.7, 2.8], [3.5, 3.6, 3.7, 3.8]]],
        dtype=np.float32)
    attention_layer = dense_attention.AdditiveAttention()
    attention_layer.build(input_shape=([1, 2, 4], [1, 3, 4]))
    # Scale tensor of shape [4]
    attention_layer.scale = np.array([[[0.5, 0.6, 0.7, 0.8]]], dtype=np.float32)
    actual = attention_layer._calculate_scores(query=q, key=k)

    # pylint:disable=line-too-long
    # expected000 = 0.5*tanh(1.+1.5) + 0.6*tanh(1.1+1.6) + 0.7*tanh(1.2+1.7) + 0.8*tanh(1.3+1.8) = 2.58044532581
    # expected001 = 0.5*tanh(1.+2.5) + 0.6*tanh(1.1+2.6) + 0.7*tanh(1.2+2.7) + 0.8*tanh(1.3+2.8) = 2.59734317449
    # expected002 = 0.5*tanh(1.+3.5) + 0.6*tanh(1.1+3.6) + 0.7*tanh(1.2+3.7) + 0.8*tanh(1.3+3.8) = 2.59964024652
    # expected010 = 0.5*tanh(2.+1.5) + 0.6*tanh(2.1+1.6) + 0.7*tanh(2.2+1.7) + 0.8*tanh(2.3+1.8) = 2.59734317449
    # expected011 = 0.5*tanh(2.+2.5) + 0.6*tanh(2.1+2.6) + 0.7*tanh(2.2+2.7) + 0.8*tanh(2.3+2.8) = 2.59964024652
    # expected012 = 0.5*tanh(2.+3.5) + 0.6*tanh(2.1+3.6) + 0.7*tanh(2.2+3.7) + 0.8*tanh(2.3+3.8) = 2.59995130916
    # pylint:enable=line-too-long
    expected = np.array([[[2.58044532581, 2.59734317449, 2.59964024652],
                          [2.59734317449, 2.59964024652, 2.59995130916]]],
                        dtype=np.float32)
    self.assertAllClose(expected, actual)

  def test_calculate_scores_one_dim_batch_size_two(self):
    # Query tensor of shape [2, 1, 1]
    q = np.array([[[1.1]], [[2.1]]], dtype=np.float32)
    # Key tensor of shape [2, 1, 1]
    k = np.array([[[1.6]], [[2.6]]], dtype=np.float32)
    attention_layer = dense_attention.AdditiveAttention()
    attention_layer.build(input_shape=([2, 1, 1], [2, 1, 1]))
    # Scale tensor of shape [1]
    attention_layer.scale = np.array([[[0.5]]], dtype=np.float32)
    actual = attention_layer._calculate_scores(query=q, key=k)

    # Expected tensor of shape [2, 1, 1].
    # expected000 = 0.5 * tanh(1.1 + 1.6) = 0.49550372683
    # expected100 = 0.5 * tanh(2.1 + 2.6) = 0.49991728277
    expected = np.array([[[0.49550372683]], [[0.49991728277]]],
                        dtype=np.float32)
    self.assertAllClose(expected, actual)

  def test_shape(self):
    # Query tensor of shape [1, 2, 4]
    q = np.array([[[1., 1.1, 1.2, 1.3], [2., 2.1, 2.2, 2.3]]], dtype=np.float32)
    # Value tensor of shape [1, 3, 4]
    v = np.array(
        [[[1.5, 1.6, 1.7, 1.8], [2.5, 2.6, 2.7, 2.8], [3.5, 3.6, 3.7, 3.8]]],
        dtype=np.float32)
    # Value mask tensor of shape [1, 3]
    v_mask = np.array([[True, True, False]], dtype=np.bool_)
    attention_layer = dense_attention.AdditiveAttention()
    actual = attention_layer([q, v], mask=[None, v_mask])

    expected_shape = [1, 2, 4]
    self.assertAllEqual(expected_shape, tf.shape(actual))

  def test_shape_no_scale(self):
    # Query tensor of shape [1, 2, 4]
    q = np.array([[[1., 1.1, 1.2, 1.3], [2., 2.1, 2.2, 2.3]]], dtype=np.float32)
    # Value tensor of shape [1, 3, 4]
    v = np.array(
        [[[1.5, 1.6, 1.7, 1.8], [2.5, 2.6, 2.7, 2.8], [3.5, 3.6, 3.7, 3.8]]],
        dtype=np.float32)
    # Value mask tensor of shape [1, 3]
    v_mask = np.array([[True, True, False]], dtype=np.bool_)
    attention_layer = dense_attention.AdditiveAttention(use_scale=False)
    actual = attention_layer([q, v], mask=[None, v_mask])

    expected_shape = [1, 2, 4]
    self.assertAllEqual(expected_shape, tf.shape(actual))

  def test_shape_with_key(self):
    # Query tensor of shape [1, 2, 4]
    q = np.array([[[1., 1.1, 1.2, 1.3], [2., 2.1, 2.2, 2.3]]], dtype=np.float32)
    # Value tensor of shape [1, 3, 4]
    v = np.array(
        [[[1.5, 1.6, 1.7, 1.8], [2.5, 2.6, 2.7, 2.8], [3.5, 3.6, 3.7, 3.8]]],
        dtype=np.float32)
    # Key tensor of shape [1, 3, 4]
    k = np.array(
        [[[1.5, 1.6, 1.7, 1.8], [2.5, 2.6, 2.7, 2.8], [3.5, 3.6, 3.7, 3.8]]],
        dtype=np.float32)
    # Value mask tensor of shape [1, 3]
    v_mask = np.array([[True, True, False]], dtype=np.bool_)
    attention_layer = dense_attention.AdditiveAttention()
    actual = attention_layer([q, v, k], mask=[None, v_mask])

    expected_shape = [1, 2, 4]
    self.assertAllEqual(expected_shape, tf.shape(actual))

  def test_multi_dim(self):
    # Query tensor of shape [1, 1, 1]
    q = np.array([[[1.1]]], dtype=np.float32)
    # Value tensor of shape [1, 3, 1]
    v = np.array([[[1.6], [0.7], [-0.8]]], dtype=np.float32)
    # Value mask tensor of shape [1, 3]
    v_mask = np.array([[True, True, False]], dtype=np.bool_)
    attention_layer = dense_attention.AdditiveAttention()
    attention_layer.build(input_shape=([1, 1, 1], [1, 3, 1]))
    # Scale tensor of shape [1]
    attention_layer.scale = np.array([[[0.5]]], dtype=np.float32)
    actual = attention_layer([q, v], mask=[None, v_mask])

    # pylint:disable=line-too-long
    # Expected scores of shape [1, 1, 3]
    # scores = [[[0.5 * tanh(1.1 + 1.6), 0.5 * tanh(1.1 + 0.7), 0.5 * tanh(1.1 - 0.8)]]]
    #        = [[[0.49550372683, 0.47340300642, 0.14565630622]]]
    # Expected attention distribution = softmax(scores) with zeros in
    # positions where v_mask == False.
    # => attention_distribution000
    #      = exp(0.49550372683)/(exp(0.49550372683) + exp(0.47340300642))
    #      = 0.50552495521
    #    attention_distribution001
    #      = exp(0.47340300642)/(exp(0.49550372683) + exp(0.47340300642))
    #      = 0.49447504478
    #    attention_distribution002 = 0
    #
    # Expected tensor of shape [1, 1, 1].
    # expected000 = 0.50552495521 * 1.6 + 0.49447504478 * 0.7 - 0 * 0.8
    #             = 1.15497245968
    # pylint:enable=line-too-long
    expected = np.array([[[1.15497245968]]], dtype=np.float32)
    self.assertAllClose(expected, actual)

  def test_multi_dim_with_key(self):
    # Query tensor of shape [1, 1, 1]
    q = np.array([[[1.1]]], dtype=np.float32)
    # Value tensor of shape [1, 3, 1]
    v = np.array([[[0.5], [0.8], [-0.3]]], dtype=np.float32)
    # Key tensor of shape [1, 3, 1]
    k = np.array([[[1.6], [0.7], [-0.8]]], dtype=np.float32)
    # Value mask tensor of shape [1, 3]
    v_mask = np.array([[True, True, False]], dtype=np.bool_)
    attention_layer = dense_attention.AdditiveAttention()
    attention_layer.build(input_shape=([1, 1, 1], [1, 3, 1]))
    # Scale tensor of shape [1]
    attention_layer.scale = np.array([[[0.5]]], dtype=np.float32)
    actual = attention_layer([q, v, k], mask=[None, v_mask])

    # pylint:disable=line-too-long
    # Expected scores of shape [1, 1, 3]
    # scores = [[[0.5 * tanh(1.1 + 1.6), 0.5 * tanh(1.1 + 0.7), 0.5 * tanh(1.1 - 0.8)]]]
    #        = [[[0.49550372683, 0.47340300642, 0.14565630622]]]
    # Expected attention distribution = softmax(scores) with zeros in
    # positions where v_mask == False.
    # => attention_distribution000
    #        = exp(0.49550372683)/(exp(0.49550372683) + exp(0.47340300642))
    #        = 0.50552495521
    #    attention_distribution001
    #        = exp(0.47340300642)/(exp(0.49550372683) + exp(0.47340300642))
    #        = 0.49447504478
    #    attention_distribution002 = 0
    #
    # Expected tensor of shape [1, 1, 1].
    # expected000 = 0.50552495521 * 0.5 + 0.49447504478 * 0.8 - 0 * 0.3
    #             = 0.64834251342
    # pylint:enable=line-too-long
    expected = np.array([[[0.64834251342]]], dtype=np.float32)
    self.assertAllClose(expected, actual)

  def test_multi_dim_with_query_mask(self):
    # Query tensor of shape [1, 2, 1]
    q = np.array([[[1.1], [-0.5]]], dtype=np.float32)
    # Value tensor of shape [1, 3, 1]
    v = np.array([[[1.6], [0.7], [-0.8]]], dtype=np.float32)
    # Query mask tensor of shape [1, 2]
    q_mask = np.array([[True, False]], dtype=np.bool_)
    # Value mask tensor of shape [1, 3]
    v_mask = np.array([[True, True, False]], dtype=np.bool_)
    attention_layer = dense_attention.AdditiveAttention()
    attention_layer.build(input_shape=([1, 1, 1], [1, 3, 1]))
    # Scale tensor of shape [1]
    attention_layer.scale = np.array([[[0.5]]], dtype=np.float32)
    actual = attention_layer([q, v], mask=[q_mask, v_mask])

    # pylint:disable=line-too-long
    # Expected scores of shape [1, 2, 3]
    # scores = [[[0.5 * tanh(1.1 + 1.6), 0.5 * tanh(1.1 + 0.7), 0.5 * tanh(1.1 - 0.8)],
    #            [0.5 * tanh(-0.5 + 1.6), 0.5 * tanh(-0.5 + 0.7), 0.5 * tanh(-0.5 - 0.8)]]]
    #        = [[[0.49550372683, 0.47340300642, 0.14565630622],
    #            [0.40024951088, 0.09868766011, -0.43086157965]]]
    # Expected attention distribution = softmax(scores) with zeros in
    # positions where v_mask == False.
    # => attention_distribution000
    #        = exp(0.49550372683)/(exp(0.49550372683) + exp(0.47340300642))
    #        = 0.50552495521
    #    attention_distribution001
    #        = exp(0.47340300642)/(exp(0.49550372683) + exp(0.47340300642))
    #        = 0.49447504478
    #    attention_distribution002 = 0
    # => attention_distribution010
    #        = exp(0.40024951088)/(exp(0.40024951088) + exp(0.09868766011))
    #        = 0.57482427975
    #    attention_distribution011
    #        = exp(0.09868766011)/(exp(0.40024951088) + exp(0.09868766011))
    #        = 0.42517572025
    #    attention_distribution012 = 0
    #
    # Expected tensor of shape [1, 2, 1] with zeros where  q_mask == False.
    # expected000 = 0.50552495521 * 1.6 + 0.49447504478 * 0.7 - 0 * 0.8
    #             = 1.15497245968
    # expected000 = 0
    # pylint:enable=line-too-long
    expected = np.array([[[1.15497245968], [0.]]], dtype=np.float32)
    self.assertAllClose(expected, actual)

  def test_serialization(self):
    # Test serialization with use_scale
    layer = dense_attention.AdditiveAttention(use_scale=True)

    config = keras.layers.serialize(layer)
    new_layer = keras.layers.deserialize(config)
    self.assertEqual(new_layer.use_scale, True)

    config = layer.get_config()
    new_layer = dense_attention.AdditiveAttention.from_config(config)
    self.assertEqual(new_layer.use_scale, True)

  @testing_utils.enable_v2_dtype_behavior
  def test_mixed_float16_policy(self):
    # Test case for GitHub issue:
    # https://github.com/tensorflow/tensorflow/issues/46064
    with policy.policy_scope('mixed_float16'):
      q = tf.cast(tf.random.uniform((2, 3, 4), seed=1), 'float16')
      v = tf.cast(tf.random.uniform((2, 3, 4), seed=2), 'float16')
      k = tf.cast(tf.random.uniform((2, 3, 4), seed=3), 'float16')
      layer = dense_attention.AdditiveAttention(causal=True)
      _ = layer([q, v, k])


@combinations.generate(combinations.combine(mode=['graph', 'eager']))
class LowerTriangularMaskTest(tf.test.TestCase, parameterized.TestCase):

  def test_square_shape(self):
    actual = dense_attention._lower_triangular_mask([3, 3])
    expected = np.array(
        [[True, False, False], [True, True, False], [True, True, True]],
        dtype=np.bool_)
    self.assertAllEqual(expected, actual)

  def test_orthogonal_shape(self):
    actual = dense_attention._lower_triangular_mask([3, 2])
    expected = np.array([[True, False], [True, True], [True, True]],
                        dtype=np.bool_)
    self.assertAllEqual(expected, actual)

  def test_three_dim(self):
    actual = dense_attention._lower_triangular_mask([1, 3, 3])
    expected = np.array(
        [[[True, False, False], [True, True, False], [True, True, True]]],
        dtype=np.bool_)
    self.assertAllEqual(expected, actual)


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