tf-models-nightly 2.17.0.dev20240617__py2.py3-none-any.whl → 2.20.0.dev20251205__py2.py3-none-any.whl

official/nlp/modeling/layers/block_sparse_attention_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.
@@ -14,6 +14,8 @@
 
 """Tests for block sparse attention layer."""
 
+import math
+
 from absl.testing import parameterized
 import numpy as np
 import tensorflow as tf, tf_keras
@@ -25,9 +27,36 @@ class BlockSparseAttentionTest(tf.test.TestCase, parameterized.TestCase):
 
   @parameterized.named_parameters(
       ("key_value_same_proj", None, None, [40, 80]),
+      ("key_value_same_proj_mqa", None, None, [40, 80], False, 1),
+      ("key_value_same_proj_multi_query_blocks", None, None, [40, 80], True),
+      (
+          "key_value_same_proj_multi_query_blocks_mqa",
+          None,
+          None,
+          [40, 80],
+          True,
+          1,
+      ),
       ("key_value_different_proj", 32, 60, [40, 60]),
+      ("key_value_different_proj_mqa", 32, 60, [40, 60], False, 1),
+      ("key_value_different_proj_multi_query_blocks", 32, 60, [40, 60], True),
+      (
+          "key_value_different_proj_multi_query_blocks_mqa",
+          32,
+          60,
+          [40, 60],
+          True,
+          1,
+      ),
   )
-  def test_non_masked_attention(self, value_dim, output_shape, output_dims):
+  def test_non_masked_attention(
+      self,
+      value_dim,
+      output_shape,
+      output_dims,
+      multi_query_blocks=False,
+      num_kv_heads=None,
+  ):
     """Test that the attention layer can be created without a mask tensor."""
     test_layer = block_sparse_attention.MultiHeadAttention(
         num_heads=12,
@@ -35,7 +64,8 @@ class BlockSparseAttentionTest(tf.test.TestCase, parameterized.TestCase):
         value_dim=value_dim,
         output_shape=output_shape,
         src_block_size=10,
-        tgt_block_size=5,
+        tgt_block_size=20 if multi_query_blocks else 5,
+        num_kv_heads=num_kv_heads,
     )
     # Create a 3-dimensional input (the first dimension is implicit).
     query = tf_keras.Input(shape=(40, 80))
@@ -53,12 +83,41 @@ class BlockSparseAttentionTest(tf.test.TestCase, parameterized.TestCase):
     output = test_layer(query, query)
     self.assertEqual(output.shape.as_list(), [None, 40, 80])
 
-  @parameterized.named_parameters(("with_bias", True), ("no_bias", False))
-  def test_masked_attention(self, use_bias):
+  @parameterized.named_parameters(
+      ("with_bias", True),
+      ("with_bias_mqa", True, False, False, 1),
+      ("with_bias_multi_query_blocks", True, False, True),
+      ("with_bias_multi_query_blocks_mqa", True, False, True, 1),
+      ("no_bias", False),
+      ("no_bias_mqa", False, False, False, 1),
+      ("no_bias_multi_query_blocks", False, False, True),
+      ("no_bias_multi_query_blocks_mqa", False, False, True, 1),
+      ("with_sigmoid_attn", True, True),
+      ("with_sigmoid_attn_mqa", True, True, False, 1),
+      ("with_sigmoid_attn_multi_query_blocks", True, True, True),
+      ("with_sigmoid_attn_multi_query_blocks_mqa", True, True, True, 1),
+  )
+  def test_masked_attention(
+      self,
+      use_bias,
+      use_sigmoid_attn=False,
+      multi_query_blocks=False,
+      num_kv_heads=None,
+  ):
     """Test with a mask tensor."""
+    if use_sigmoid_attn:
+      sigmoid_attn_bias = -math.log(2)
+    else:
+      sigmoid_attn_bias = None
     test_layer = block_sparse_attention.MultiHeadAttention(
-        num_heads=4, key_dim=2, use_bias=use_bias, src_block_size=2,
-        tgt_block_size=1,
+        num_heads=4,
+        key_dim=2,
+        use_bias=use_bias,
+        src_block_size=2,
+        tgt_block_size=2 if multi_query_blocks else 1,
+        use_sigmoid_attn=use_sigmoid_attn,
+        sigmoid_attn_bias=sigmoid_attn_bias,
+        num_kv_heads=num_kv_heads,
     )
     # Create a 3-dimensional input (the first dimension is implicit).
     batch_size = 3
@@ -112,6 +171,77 @@ class BlockSparseAttentionTest(tf.test.TestCase, parameterized.TestCase):
       self.assertLen(test_layer._query_dense.trainable_variables, 1)
       self.assertLen(test_layer._output_dense.trainable_variables, 1)
 
+  @parameterized.named_parameters(
+      ("default_with_softmax", False),
+      ("default_with_sigmoid", True),
+  )
+  def test_default_masked_attention(
+      self,
+      use_sigmoid_attn=False,
+  ):
+    """Test with a mask tensor."""
+    seq_len = 8
+    if use_sigmoid_attn:
+      sigmoid_attn_bias = -math.log(seq_len)
+    else:
+      sigmoid_attn_bias = None
+    test_layer = block_sparse_attention.MultiHeadAttention(
+        num_heads=4,
+        key_dim=2,
+        use_bias=True,
+        src_block_size=seq_len,
+        tgt_block_size=seq_len,
+        use_sigmoid_attn=use_sigmoid_attn,
+        sigmoid_attn_bias=sigmoid_attn_bias,
+    )
+    # Create a 3-dimensional input (the first dimension is implicit).
+    batch_size = 3
+    query = tf_keras.Input(shape=(seq_len, 8))
+    value = tf_keras.Input(shape=(seq_len, 8))
+    mask_tensor = tf_keras.Input(shape=(seq_len, seq_len))
+    output = test_layer(query=query, value=value, attention_mask=mask_tensor)
+
+    # Create a model containing the test layer.
+    model = tf_keras.Model([query, value, mask_tensor], output)
+
+    # Generate data for the input (non-mask) tensors.
+    from_data = 10 * np.random.random_sample((batch_size, seq_len, 8))
+    to_data = 10 * np.random.random_sample((batch_size, seq_len, 8))
+
+    # Invoke the data with a random set of mask data. This should mask at
+    # least one element.
+    mask_data = np.random.randint(2, size=(batch_size, seq_len, seq_len))
+    masked_output_data = model.predict([from_data, to_data, mask_data])
+
+    # Invoke the same data, but with a null mask (where no elements are
+    # masked).
+    null_mask_data = np.ones((batch_size, seq_len, seq_len))
+    unmasked_output_data = model.predict([from_data, to_data, null_mask_data])
+
+    # Because one data is masked and one is not, the outputs should not be
+    # the same.
+    self.assertNotAllClose(masked_output_data, unmasked_output_data)
+
+    # Tests the layer with three inputs: Q, K, V.
+    key = tf_keras.Input(shape=(seq_len, 8))
+    output = test_layer(
+        query, value=value, key=key, attention_mask=mask_tensor
+    )
+    model = tf_keras.Model([query, value, key, mask_tensor], output)
+
+    masked_output_data = model.predict(
+        [from_data, to_data, to_data, mask_data]
+    )
+    unmasked_output_data = model.predict(
+        [from_data, to_data, to_data, null_mask_data]
+    )
+    # Because one data is masked and one is not, the outputs should not be
+    # the same.
+    self.assertNotAllClose(masked_output_data, unmasked_output_data)
+
+    self.assertLen(test_layer._query_dense.trainable_variables, 2)
+    self.assertLen(test_layer._output_dense.trainable_variables, 2)
+
   def test_masked_attention_with_scores(self):
     """Test with a mask tensor."""
     test_layer = block_sparse_attention.MultiHeadAttention(

official/nlp/modeling/layers/cls_head.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/nlp/modeling/layers/cls_head_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/nlp/modeling/layers/factorized_embedding.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/nlp/modeling/layers/factorized_embedding_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/nlp/modeling/layers/gated_feedforward.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.
@@ -37,7 +37,7 @@ class GatedFeedforward(tf_keras.layers.Layer):
     dropout: Dropout probability for the output dropout.
     use_gate: Whether to use gated linear units. If True, assuming `GELU` as the
       activation and omitting bias, will apply
-      `GEGLU(x, W, V, W_2) = (GEGLU(xW) * xV)W2`; if False, will follow
+      `GEGLU(x, W, V, W_2) = (GELU(xW) * xV)W2`; if False, will follow
       "Attention Is All You Need" (https://arxiv.org/abs/1706.03762) paper and
         apply `FFN(x, W, W_2) = GELU(xW_1)W_2.`
     num_blocks: The number of feedforward blocks to stack. Each block contains a

official/nlp/modeling/layers/gated_feedforward_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/nlp/modeling/layers/gaussian_process.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/nlp/modeling/layers/gaussian_process_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/nlp/modeling/layers/kernel_attention.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/nlp/modeling/layers/kernel_attention_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/nlp/modeling/layers/masked_lm.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/nlp/modeling/layers/masked_lm_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/nlp/modeling/layers/masked_softmax.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/nlp/modeling/layers/masked_softmax_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/nlp/modeling/layers/mat_mul_with_margin.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.
@@ -16,7 +16,6 @@
 # pylint: disable=g-classes-have-attributes
 
 from typing import Tuple
-# Import libraries
 import tensorflow as tf, tf_keras
 
 from official.modeling import tf_utils

official/nlp/modeling/layers/mat_mul_with_margin_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/nlp/modeling/layers/mixing.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/nlp/modeling/layers/mixing_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/nlp/modeling/layers/mobile_bert_layers.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/nlp/modeling/layers/mobile_bert_layers_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/nlp/modeling/layers/moe.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/nlp/modeling/layers/moe_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/nlp/modeling/layers/multi_channel_attention.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/nlp/modeling/layers/multi_channel_attention_test.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.

official/nlp/modeling/layers/multi_query_attention.py

@@ -1,4 +1,4 @@
-# Copyright 2024 The TensorFlow Authors. All Rights Reserved.
+# Copyright 2025 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.
@@ -18,10 +18,13 @@ Based on https://arxiv.org/pdf/1911.02150.pdf and
 https://arxiv.org/pdf/2305.13245.pdf.
 """
 
+import math
 import string
 from typing import Optional, Sequence, Union
 
+import gin
 import tensorflow as tf, tf_keras
+from official.modeling import tf_utils
 
 _CHR_IDX = string.ascii_lowercase
 
@@ -78,7 +81,9 @@ def _get_output_shape(
 class MultiHeadAttention(tf_keras.layers.MultiHeadAttention):
   """Multi-query attention layer."""
 
-  def __init__(self, num_kv_heads=None, **kwargs):
+  def __init__(
+      self, num_kv_heads=None, enable_gqa_optimization=False, **kwargs
+  ):
     # num_kv_heads defines the number of key/value heads. A value of 1 means
     # that the key/value heads are shared across all query heads. Any other
     # value must be less than num_heads and must divide num_heads exactly. If
@@ -86,6 +91,16 @@ class MultiHeadAttention(tf_keras.layers.MultiHeadAttention):
     # num_kv_heads.
     super().__init__(**kwargs)
     self._num_kv_heads = num_kv_heads or self._num_heads
+    # TODO(akandoor): Remove this flag once the GQA optimization is rolled out.
+    # This flag is used to enable order of K,G in the einsum equations.
+    # This optimization is only used in GQA, and is disabled by default.
+    # If enabled, the einsum equations are:
+    #   1. Dot product: "...SKH,...TKGH->...KGTS"
+    #   2. Combine: "...KGTS,...SKH->...TKGH"
+    # If disabled, the einsum equations are:
+    #   1. Dot product: "...SKH,...TKnH->...nKTS"
+    #   2. Combine: "...nKTS,...SKH->...TnKH"
+    self._enable_gqa_optimization = enable_gqa_optimization
     assert (
         self._num_kv_heads < self._num_heads
     ), "num_kv_heads must be less than num_heads."
@@ -93,6 +108,11 @@ class MultiHeadAttention(tf_keras.layers.MultiHeadAttention):
         self._num_heads % self._num_kv_heads == 0
     ), "num_kv_heads needs to divide num_heads exactly."
 
+  def get_config(self):
+    config = super().get_config()
+    config.update({"num_kv_heads": self._num_kv_heads})
+    return config
+
   def _build_from_signature(
       self,
       query: Union[tf.Tensor, tf.TensorShape],
@@ -138,8 +158,12 @@ class MultiHeadAttention(tf_keras.layers.MultiHeadAttention):
         output_dims = 2
         key_last_dims = [self._num_kv_heads, self._key_dim]
         value_last_dims = [self._num_kv_heads, self._value_dim]
-        self._dot_product_equation = "...SKH,...TKnH->...nKTS"
-        self._combine_equation = "...nKTS,...SKH->...TnKH"
+        if self._enable_gqa_optimization:
+          self._dot_product_equation = "...SKH,...TKGH->...KGTS"
+          self._combine_equation = "...KGTS,...SKH->...TKGH"
+        else:
+          self._dot_product_equation = "...SKH,...TKnH->...nKTS"
+          self._combine_equation = "...nKTS,...SKH->...TnKH"
 
       einsum_equation, bias_axes, output_rank = _build_proj_equation(
           free_dims=self._key_shape.rank - 1,
@@ -165,6 +189,9 @@ class MultiHeadAttention(tf_keras.layers.MultiHeadAttention):
           name="value",
           **self._get_common_kwargs_for_sublayer(),
       )
+      self._qkv_rank = (
+          output_rank if self._num_kv_heads > 1 else output_rank + 1
+      )
 
   def _compute_attention(
       self, query, key, value, attention_mask=None, training=None
@@ -206,3 +233,194 @@ class MultiHeadAttention(tf_keras.layers.MultiHeadAttention):
           ],
       )
     return attention_output, attention_scores
+
+
+@tf_keras.utils.register_keras_serializable(package="Text")
+@gin.configurable
+class TalkingHeadsMultiQueryAttention(MultiHeadAttention):
+  """Implements Talking-Heads Attention combined with Multi-Query Attention.
+
+  See https://arxiv.org/pdf/2003.02436 for more details.
+  TODO(akandoor): Make num talking heads configurable. Currently, num talking
+  heads is fixed to num query heads.
+
+  This class inherits from MultiQueryAttention to get the MQA-specific
+  logic for __init__, get_config.
+
+  It then overrides _build_from_signature to add the talking-heads weights
+  and overrides _compute_attention to merge the MQA wrapper (for
+  reshaping) with the THA computation (for pre/post-softmax projections).
+  """
+
+  def _build_from_signature(
+      self,
+      query: Union[tf.Tensor, tf.TensorShape],
+      value: Union[tf.Tensor, tf.TensorShape],
+      key: Optional[Union[tf.Tensor, tf.TensorShape]] = None,
+  ):
+    """Builds layers and variables."""
+    # Call the parent (MultiQueryAttention) _build_from_signature.
+    super()._build_from_signature(query=query, value=value, key=key)
+    # Now, *after* all MQA setup is done, we add the THA setup logic.
+    qkv_rank = self._qkv_rank
+    # TalkingHeadsAttention logic to the MQA build logic.
+    num_batch_dims = qkv_rank - len(self._attention_axes) - 2
+    attn_scores_rank = num_batch_dims + 1 + len(self._attention_axes) * 2
+    scores_notation = _CHR_IDX[:attn_scores_rank]
+    projection_notation = scores_notation[num_batch_dims] + (
+        _CHR_IDX[attn_scores_rank])
+    projected_scores_notation = scores_notation[:num_batch_dims] + (
+        _CHR_IDX[attn_scores_rank] + scores_notation[num_batch_dims + 1:])
+    self._talking_heads_equation = "%s,%s->%s" % (
+        scores_notation, projection_notation, projected_scores_notation)
+
+    with tf.init_scope():
+      self._pre_softmax_weight = self.add_weight(
+          "pre_softmax_weight",
+          shape=(self._num_heads, self._num_heads),
+          initializer=tf_utils.clone_initializer(self._kernel_initializer),
+          regularizer=self._kernel_regularizer,
+          constraint=self._kernel_constraint,
+          dtype=self.dtype,
+          trainable=True)
+      self._post_softmax_weight = self.add_weight(
+          "post_softmax_weight",
+          shape=(self._num_heads, self._num_heads),
+          initializer=tf_utils.clone_initializer(self._kernel_initializer),
+          regularizer=self._kernel_regularizer,
+          constraint=self._kernel_constraint,
+          dtype=self.dtype,
+          trainable=True)
+
+  def _compute_attention(
+      self, query, key, value, attention_mask=None, training=None
+  ):
+    """Applies Dot-product attention, merging MQA wrapper and THA computation.
+
+    Args:
+      query: Projected query `Tensor` of shape `[B, T, N, key_dim]`.
+      key: Projected key `Tensor` of shape `[B, T, N, key_dim]`.
+      value: Projected value `Tensor` of shape `[B, T, N, value_dim]`.
+      attention_mask: a boolean mask of shape `[B, T, S]`, that prevents
+        attention to certain positions.
+      training: Python boolean indicating whether the layer should behave in
+        training mode (adding dropout) or in inference mode (doing nothing).
+
+    Returns:
+      attention_output: Multi-headed outputs of attention computation.
+      attention_scores: Multi-headed attention weights.
+    """
+    # This is the MQA "wrapper" logic for grouped queries
+    query_shape = tf.shape(query)
+    if self._num_kv_heads > 1:
+      query = tf.reshape(
+          query,
+          [
+              query_shape[0],
+              query_shape[1],
+              self._num_kv_heads,
+              self._num_heads // self._num_kv_heads,
+              query_shape[-1],
+          ],
+      )
+
+    # This is the THA "computation" logic
+    # Note: Applying scalar multiply at the smaller end of einsum improves
+    # XLA performance, but may introduce slight numeric differences in
+    # the Transformer attention head.
+    query = tf.multiply(
+        query, 1.0 / math.sqrt(float(self._key_dim))
+    )
+
+    # Note: self._dot_product_equation was set by _build_from_signature
+    # (from MQA) to be MQA-compatible.
+    attention_scores = tf.einsum(self._dot_product_equation, key, query)
+
+    # --- Talking-Heads modification for MQA ---
+    # The THA _talking_heads_equation expects scores of shape [B, N, T, S].
+    # The MQA _dot_product_equation produces [B, K, G, T, S].
+    # We must reshape before and after applying TH logic.
+    scores_shape = tf.shape(attention_scores)
+    if self._num_kv_heads > 1:
+      # Reshape from [B, K, G, T, S] to [B, N, T, S]
+      attention_scores = tf.reshape(
+          attention_scores,
+          [
+              scores_shape[0],  # Batch
+              self._num_heads,  # N = K * G
+              scores_shape[-2],  # T
+              scores_shape[-1]   # S
+          ]
+      )
+
+    # Apply linear projection before softmax
+    attention_scores = tf.einsum(self._talking_heads_equation, attention_scores,
+                                 self._pre_softmax_weight)
+
+    # Normalize the attention scores to probabilities.
+    attention_scores = self._masked_softmax(attention_scores, attention_mask)
+
+    # Apply linear projection after softmax
+    attention_scores = tf.einsum(self._talking_heads_equation, attention_scores,
+                                 self._post_softmax_weight)
+
+    # Reshape back to MQA-compatible shape [B, K, G, T, S]
+    # before the final combine_equation
+    if self._num_kv_heads > 1:
+      if self._enable_gqa_optimization:
+        attention_scores = tf.reshape(
+            attention_scores,
+            [
+                scores_shape[0],  # B
+                self._num_kv_heads,  # K
+                self._num_heads // self._num_kv_heads,  # G
+                scores_shape[-2],  # T
+                scores_shape[-1]   # S
+            ]
+        )
+      else:
+        attention_scores = tf.reshape(
+            attention_scores,
+            [
+                scores_shape[0],  # B
+                self._num_heads // self._num_kv_heads,  # G
+                self._num_kv_heads,  # K
+                scores_shape[-2],  # T
+                scores_shape[-1]   # S
+            ]
+        )
+
+    # This is actually dropping out entire tokens to attend to.
+    attention_scores_dropout = self._dropout_layer(
+        attention_scores, training=training)
+
+    # Note: self._combine_equation was set by _build_from_signature
+    # (from MQA) to be MQA-compatible.
+    attention_output = tf.einsum(self._combine_equation,
+                                 attention_scores_dropout, value)
+
+    # This is the MQA "wrapper" logic for grouped queries
+    if self._num_kv_heads > 1:
+      attention_output = tf.reshape(
+          attention_output,
+          [
+              query_shape[0],
+              query_shape[1],
+              self._num_heads,
+              tf.shape(attention_output)[-1],
+          ],
+      )
+      # We also need to reshape the final scores back to [B, N, T, S]
+      # for the return value.
+      attention_scores = tf.reshape(
+          attention_scores,
+          [
+              query_shape[0],
+              self._num_heads,
+              tf.shape(attention_scores)[-2],
+              tf.shape(attention_scores)[-1],
+          ],
+      )
+
+    return attention_output, attention_scores
+