text: # 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.
  
  """Roformer attention layer."""
  # pylint: disable=g-classes-have-attributes
  import tensorflow as tf
  
  EinsumDense = tf.keras.layers.experimental.EinsumDense
  MultiHeadAttention = tf.keras.layers.MultiHeadAttention
  
  
  def _build_trig_vector(length, key_dim):
    """Builds the trig vector."""
    tf_dtype = tf.keras.mixed_precision.global_policy().compute_dtype
    position_ids = tf.cast(tf.range(length), dtype=tf_dtype)
    position_ids = tf.expand_dims(position_ids, axis=0)
    steps = key_dim // 2
    indices = tf.cast(tf.range(steps), dtype=tf_dtype)
    indices = tf.pow(tf.constant(10000.0, dtype=tf_dtype), -2 * indices / steps)
    vec = tf.einsum('bl,d->bld', position_ids, indices)
    sin_vec = tf.repeat(tf.sin(vec), repeats=2, axis=-1)
    cos_vec = tf.repeat(tf.cos(vec), repeats=2, axis=-1)
    sin_vec, cos_vec = tf.expand_dims(sin_vec, 2), tf.expand_dims(cos_vec, 2)
    return sin_vec, cos_vec
  
  
  @tf.keras.utils.register_keras_serializable(package='Text')
  class RoformerAttention(tf.keras.layers.MultiHeadAttention):
    """Roformer Attention."""
  
    def __init__(self,
                 q_max_sequence_length,
                 kv_max_sequence_length,
                 output_range=None,
                 **kwargs):
      """Instantiates a roformer attention layer.
  
      Roformer paper: https://arxiv.org/abs/2104.09864
  
      Args:
        q_max_sequence_length: maximum length in input for the query
        kv_max_sequence_length: maximum length in input for key and value, can be
          different from q_max_sequence_length
        output_range: length of the query tensor to consider.
        **kwargs: other keyword arguments.
      """
      super().__init__(**kwargs)
      self._q_max_sequence_length = q_max_sequence_length
      self._kv_max_sequence_length = kv_max_sequence_length
      assert self._key_dim % 2 == 0
      q_sin_vec, q_cos_vec = _build_trig_vector(self._q_max_sequence_length,
                                                self._key_dim)
      k_sin_vec, k_cos_vec = _build_trig_vector(self._kv_max_sequence_length,
                                                self._key_dim)
      # pylint:disable=g-long-ternary
      self.q_sin_vec, self.q_cos_vec = (q_sin_vec,
                                        q_cos_vec) if output_range is None else (
                                            q_sin_vec[:, 0:output_range, ...],
                                            q_cos_vec[:, 0:output_range, ...])
      # pylint:enable=g-long-ternary
      self.k_sin_vec, self.k_cos_vec = (k_sin_vec, k_cos_vec)
  
    def roformer_recompute_qkv(self, q, k, v):
      q_shape = tf.shape(q)
      q_len = q_shape[1]
      k_shape = tf.shape(k)
      k_len = k_shape[1]
  
      q2 = tf.stack([-q[..., 1::2], q[..., ::2]], axis=4)
      q2 = tf.reshape(q2, q_shape)
      k2 = tf.stack([-k[..., 1::2], k[..., ::2]], axis=4)
      k2 = tf.reshape(k2, k_shape)
      ret_q = q * self.q_cos_vec[:, 0:q_len,
                                 ...] + q2 * self.q_sin_vec[:, 0:q_len, ...]
      ret_w = k * self.k_cos_vec[:, 0:k_len,
                                 ...] + k2 * self.k_sin_vec[:, 0:k_len, ...]
      return ret_q, ret_w, v
  
    def call(self,
             query,
             value,
             key=None,
             attention_mask=None,
             return_attention_scores=False,
             training=None):
      if not self._built_from_signature:
        self._build_from_signature(query=query, value=value, key=key)
      if key is None:
        key = value
  
      query = self._query_dense(query)
      key = self._key_dense(key)
      value = self._value_dense(value)
  
      query, key, value = self.roformer_recompute_qkv(query, key, value)
  
      attention_output, attention_scores = self._compute_attention(
          query, key, value, attention_mask, training)
      attention_output = self._output_dense(attention_output)
  
      if return_attention_scores:
        return attention_output, attention_scores
      return attention_output
  