tf-models-nightly 2.20.0.dev20251028__py2.py3-none-any.whl → 2.20.0.dev20251030__py2.py3-none-any.whl

official/nlp/modeling/layers/multi_query_attention.py

@@ -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."
@@ -143,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,
@@ -170,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
@@ -211,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
+

official/nlp/modeling/layers/multi_query_attention_test.py

@@ -211,5 +211,205 @@ class MultiQueryAttentionTest(tf.test.TestCase, parameterized.TestCase):
     self.assertNotAllClose(masked_score, unmasked_score)
 
 
+class TalkingHeadsMultiQueryAttentionTest(
+    tf.test.TestCase, parameterized.TestCase
+):
+
+  @parameterized.named_parameters(
+      ("key_value_same_proj_mqa", 1, None, None, [40, 80]),
+      ("key_value_different_proj_mqa", 1, 32, 60, [40, 60]),
+      ("key_value_same_proj_gqa", 3, None, None, [40, 80]),
+      ("key_value_different_proj_gqa", 3, 32, 60, [40, 60]),
+  )
+  def test_non_masked_attention(
+      self, num_kv_heads, value_dim, output_shape, output_dims
+  ):
+    """Test that the attention layer can be created without a mask tensor."""
+    test_layer = multi_query_attention.TalkingHeadsMultiQueryAttention(
+        num_heads=12,
+        num_kv_heads=num_kv_heads,
+        enable_gqa_optimization=num_kv_heads > 1,
+        key_dim=64,
+        value_dim=value_dim,
+        output_shape=output_shape,
+    )
+    # Create a 3-dimensional input (the first dimension is implicit).
+    query = tf_keras.Input(shape=(40, 80))
+    value = tf_keras.Input(shape=(20, 80))
+    output = test_layer(query=query, value=value)
+    self.assertEqual(output.shape.as_list(), [None] + output_dims)
+
+  @parameterized.named_parameters(
+      ("_mqa", 1),
+      ("_gqa", 3),
+  )
+  def test_non_masked_self_attention(self, num_kv_heads):
+    """Test with one input (self-attenntion) and no mask tensor."""
+    test_layer = multi_query_attention.TalkingHeadsMultiQueryAttention(
+        num_heads=12,
+        num_kv_heads=num_kv_heads,
+        enable_gqa_optimization=num_kv_heads > 1,
+        key_dim=64,
+    )
+    # Create a 3-dimensional input (the first dimension is implicit).
+    query = tf_keras.Input(shape=(40, 80))
+    output = test_layer(query, query)
+    self.assertEqual(output.shape.as_list(), [None, 40, 80])
+
+  @parameterized.named_parameters(
+      ("_mqa", 1),
+      ("_gqa", 3),
+  )
+  def test_attention_scores(self, num_kv_heads):
+    """Test attention outputs with coefficients."""
+    test_layer = multi_query_attention.TalkingHeadsMultiQueryAttention(
+        num_heads=12, num_kv_heads=num_kv_heads, key_dim=64
+    )
+    # Create a 3-dimensional input (the first dimension is implicit).
+    query = tf_keras.Input(shape=(40, 80))
+    output, coef = test_layer(query, query, return_attention_scores=True)
+    self.assertEqual(output.shape.as_list(), [None, 40, 80])
+    self.assertEqual(coef.shape.as_list(), [None, 12, 40, 40])
+
+  @parameterized.named_parameters(
+      ("_mqa", 1),
+      ("_gqa", 3),
+  )
+  def test_attention_scores_with_values(self, num_kv_heads):
+    """Test attention outputs with coefficients."""
+    test_layer = multi_query_attention.TalkingHeadsMultiQueryAttention(
+        num_heads=12, num_kv_heads=num_kv_heads, key_dim=64
+    )
+    # Create a 3-dimensional input (the first dimension is implicit).
+    query = tf_keras.Input(shape=(40, 80))
+    value = tf_keras.Input(shape=(60, 80))
+    output, coef = test_layer(query, value, return_attention_scores=True)
+    self.assertEqual(output.shape.as_list(), [None, 40, 80])
+    self.assertEqual(coef.shape.as_list(), [None, 12, 40, 60])
+
+  @parameterized.named_parameters(
+      ("with_bias_mqa", 1, True),
+      ("no_bias_mqa", 1, False),
+      ("with_bias_gqa", 2, True),
+      ("no_bias_gqa", 2, False),
+  )
+  def test_masked_attention(self, num_kv_heads, use_bias):
+    """Test with a mask tensor."""
+    test_layer = multi_query_attention.TalkingHeadsMultiQueryAttention(
+        num_heads=4,
+        num_kv_heads=num_kv_heads,
+        enable_gqa_optimization=num_kv_heads > 1,
+        key_dim=2,
+        use_bias=use_bias,
+    )
+    # Create a 3-dimensional input (the first dimension is implicit).
+    batch_size = 3
+    query = tf_keras.Input(shape=(4, 8))
+    value = tf_keras.Input(shape=(2, 8))
+    mask_tensor = tf_keras.Input(shape=(4, 2))
+    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, 4, 8))
+    to_data = 10 * np.random.random_sample((batch_size, 2, 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, 4, 2))
+    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, 4, 2))
+    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=(2, 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)
+
+    if use_bias:
+      self.assertLen(test_layer._query_dense.trainable_variables, 2)
+      self.assertLen(test_layer._output_dense.trainable_variables, 2)
+    else:
+      self.assertLen(test_layer._query_dense.trainable_variables, 1)
+      self.assertLen(test_layer._output_dense.trainable_variables, 1)
+
+  @parameterized.named_parameters(
+      ("_mqa", 1),
+      ("_gqa", 2),
+  )
+  def test_masked_attention_with_scores(self, num_kv_heads):
+    """Test with a mask tensor."""
+    test_layer = multi_query_attention.TalkingHeadsMultiQueryAttention(
+        num_heads=4, num_kv_heads=num_kv_heads, key_dim=2
+    )
+    # Create a 3-dimensional input (the first dimension is implicit).
+    batch_size = 3
+    query = tf_keras.Input(shape=(4, 8))
+    value = tf_keras.Input(shape=(2, 8))
+    mask_tensor = tf_keras.Input(shape=(4, 2))
+    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, 4, 8))
+    to_data = 10 * np.random.random_sample((batch_size, 2, 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, 4, 2))
+    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, 4, 2))
+    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)
+
+    # Create a model containing attention scores.
+    output, scores = test_layer(
+        query=query,
+        value=value,
+        attention_mask=mask_tensor,
+        return_attention_scores=True,
+    )
+    model = tf_keras.Model([query, value, mask_tensor], [output, scores])
+    masked_output_data_score, masked_score = model.predict(
+        [from_data, to_data, mask_data]
+    )
+    unmasked_output_data_score, unmasked_score = model.predict(
+        [from_data, to_data, null_mask_data]
+    )
+    self.assertNotAllClose(masked_output_data_score, unmasked_output_data_score)
+    self.assertAllClose(masked_output_data, masked_output_data_score)
+    self.assertAllClose(unmasked_output_data, unmasked_output_data_score)
+    self.assertNotAllClose(masked_score, unmasked_score)
+
+
 if __name__ == "__main__":
   tf.test.main()

official/nlp/modeling/layers/transformer_encoder_block.py

@@ -20,6 +20,7 @@ import tensorflow as tf, tf_keras
 from official.modeling import tf_utils
 from official.nlp.modeling.layers import block_sparse_attention
 from official.nlp.modeling.layers import multi_query_attention
+from official.nlp.modeling.layers import talking_heads_attention
 from official.nlp.modeling.layers import util
 
 
@@ -118,6 +119,8 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
       linformer_dim=None,
       linformer_shared_kv_projection=True,
       lowrank_query_seq_proj_dim=None,
+      enable_talking_heads=False,
+      enable_gqa_optimization=False,
       **kwargs,
   ):
     """Initializes `TransformerEncoderBlock`.
@@ -202,6 +205,10 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
         keys and values.
       lowrank_query_seq_proj_dim: If set, applies a projection layer on query
         sequence to the given dimension. go/constformer-doc
+      enable_talking_heads: Enable talking heads as in
+        https://arxiv.org/pdf/2003.02436.
+      enable_gqa_optimization: Enable GQA optimization in multi-query attention.
+        This flag is valid only when num_kv_heads is set for GQA.
       **kwargs: keyword arguments.
     """
     util.filter_kwargs(kwargs)
@@ -244,6 +251,8 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
     self._linformer_dim = linformer_dim
     self._linformer_shared_kv_projection = linformer_shared_kv_projection
     self._lowrank_query_seq_proj_dim = lowrank_query_seq_proj_dim
+    self._enable_talking_heads = enable_talking_heads
+    self._enable_gqa_optimization = enable_gqa_optimization
     if (
         self._src_block_size is not None
         and self._num_kv_heads is not None
@@ -314,6 +323,11 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
         bias_constraint=self._bias_constraint,
     )
     if self._src_block_size is not None:
+      if self._enable_talking_heads:
+        raise ValueError(
+            "Block sparse attention does not support talking heads. Please"
+            " set enable_talking_heads to False."
+        )
       attention_layer_kwargs.update(
           src_block_size=self._src_block_size,
           tgt_block_size=self._tgt_block_size,
@@ -326,9 +340,22 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
     elif self._num_kv_heads is not None:
       attention_layer_kwargs.update(
           num_kv_heads=self._num_kv_heads,
+          enable_gqa_optimization=self._enable_gqa_optimization,
           name="multi_query_attention",
       )
-      attention_fn = multi_query_attention.MultiHeadAttention
+      if self._enable_talking_heads:
+        attention_fn = (
+            multi_query_attention.TalkingHeadsMultiQueryAttention
+        )
+      else:
+        attention_fn = multi_query_attention.MultiHeadAttention
+    elif self._enable_talking_heads:
+      attention_layer_kwargs.update(
+          name="talking_heads_attention",
+      )
+      attention_fn = (
+          talking_heads_attention.TalkingHeadsAttention
+      )
     else:
       attention_fn = tf_keras.layers.MultiHeadAttention
     self._attention_layer = attention_fn(

official/nlp/modeling/layers/transformer_encoder_block_test.py

@@ -743,35 +743,46 @@ class TransformerArgumentTest(tf.test.TestCase, parameterized.TestCase):
         num_attention_heads=num_attention_heads,
         inner_dim=2048,
         inner_activation='relu',
-        return_attention_scores=return_attention_scores)
+        return_attention_scores=return_attention_scores,
+    )
     # Create a 3-dimensional input (the first dimension is implicit).
     data_tensor = tf_keras.Input(shape=(sequence_length, width))
     output_tensor = test_layer(data_tensor)
 
     expected_layer_output_shape = [None, sequence_length, width]
     expected_attention_scores_shape = [
-        None, num_attention_heads, sequence_length, sequence_length
+        None,
+        num_attention_heads,
+        sequence_length,
+        sequence_length,
     ]
 
     if return_attention_scores:
       self.assertIsInstance(output_tensor, tuple)
       self.assertLen(output_tensor, 2)
       # First is the standard output.
-      self.assertEqual(output_tensor[0].shape.as_list(),
-                       expected_layer_output_shape)
+      self.assertEqual(
+          output_tensor[0].shape.as_list(), expected_layer_output_shape
+      )
       # Second is the attention scores.
-      self.assertEqual(output_tensor[1].shape.as_list(),
-                       expected_attention_scores_shape)
+      self.assertEqual(
+          output_tensor[1].shape.as_list(), expected_attention_scores_shape
+      )
     else:
       # Only the standard layer output.
-      self.assertEqual(output_tensor.shape.as_list(),
-                       expected_layer_output_shape)
+      self.assertEqual(
+          output_tensor.shape.as_list(), expected_layer_output_shape
+      )
 
   @parameterized.named_parameters(
       ('mqa', 1),
       ('gqa', 4),
+      ('talking_heads_mqa', 1, True),
+      ('talking_heads_gqa', 4, True),
   )
-  def test_attention_with_kv_heads(self, num_kv_heads):
+  def test_attention_with_kv_heads(
+      self, num_kv_heads, enable_talking_heads=False
+  ):
     num_attention_heads = 8
     sequence_length = 21
     width = 80
@@ -782,6 +793,8 @@ class TransformerArgumentTest(tf.test.TestCase, parameterized.TestCase):
         inner_activation='relu',
         return_attention_scores=True,
         num_kv_heads=num_kv_heads,
+        enable_talking_heads=enable_talking_heads,
+        enable_gqa_optimization=True,
     )
     # Create a 3-dimensional input (the first dimension is implicit).
     data_tensor = tf_keras.Input(shape=(sequence_length, width))

{tf_models_nightly-2.20.0.dev20251028.dist-info → tf_models_nightly-2.20.0.dev20251030.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.1
 Name: tf-models-nightly
-Version: 2.20.0.dev20251028
+Version: 2.20.0.dev20251030
 Summary: TensorFlow Official Models
 Home-page: https://github.com/tensorflow/models
 Author: Google Inc.

{tf_models_nightly-2.20.0.dev20251028.dist-info → tf_models_nightly-2.20.0.dev20251030.dist-info}/RECORD

@@ -331,8 +331,8 @@ official/nlp/modeling/layers/moe.py,sha256=vG5jNwX2Vt3GjWDZ_viOCfgI77antXDAnBtho
 official/nlp/modeling/layers/moe_test.py,sha256=NmLln8M820Z4VX9BYFmBppiv8_NKoTWWzNGZ-g6eEBQ,9414
 official/nlp/modeling/layers/multi_channel_attention.py,sha256=_9py6IXTJpBM3lhK0Xh4zsmrO7S8xbQJ4CR5PcKaPk4,7322
 official/nlp/modeling/layers/multi_channel_attention_test.py,sha256=X-kLm-w_jbITZgGNSKAjEPOtjH2xJABI-mavBn2jbtA,1922
-official/nlp/modeling/layers/multi_query_attention.py,sha256=Tsv1tg3GMcze9cpCpNddamL4dHWFQr4uyqO4n4cAJJk,7111
-official/nlp/modeling/layers/multi_query_attention_test.py,sha256=HplLhNbhr24r3wBD5Dzz7Y_VxylQ5DEa1prmXYmSN3g,8512
+official/nlp/modeling/layers/multi_query_attention.py,sha256=NvFBSiGZMdQnikRMSkG3kKX5e7nP2qvUt0nYFTajMqc,15384
+official/nlp/modeling/layers/multi_query_attention_test.py,sha256=2KLIynj9rXOcuvMxQ7s1JWo0B_C31RhBrsxGLX4ijg8,16455
 official/nlp/modeling/layers/on_device_embedding.py,sha256=RN01ud3FeInON_bXVtzItfTLivolbu45R6ajPp1_BIo,4582
 official/nlp/modeling/layers/on_device_embedding_test.py,sha256=i9esQYblo_uC-UnSkRs8Xgr6BldZ48iKoZhN4vrSTXI,8589
 official/nlp/modeling/layers/pack_optimization.py,sha256=0CoWos2h4dfqO3UvI2xLpHnWzvJpDiBJVHR2MTamLM4,10760
@@ -363,8 +363,8 @@ official/nlp/modeling/layers/tn_expand_condense_test.py,sha256=QWq1dJqQUPe5n69K3
 official/nlp/modeling/layers/tn_transformer_expand_condense.py,sha256=omzTkCBEk2TOkHEYDEBwve6WsOitX7IIJHzeKXdqDq0,11012
 official/nlp/modeling/layers/tn_transformer_test.py,sha256=pSCONEZRI4J9_6QLTJ3g_ynUYLrRXsJ1c2YMSiOV_60,8893
 official/nlp/modeling/layers/transformer.py,sha256=VjUO-gVj_PnavbT_vSrg5NDKMr0SRSiqSg5ktd42m5M,20087
-official/nlp/modeling/layers/transformer_encoder_block.py,sha256=5GJgtK1mdTxMDYVWfUoBAI_GvjDL0zO9AWtKCovSZiU,28789
-official/nlp/modeling/layers/transformer_encoder_block_test.py,sha256=7yBgv1UNmfOFre6txF_Rq93RLc1TJwnJ7-Dz4p55sy4,37602
+official/nlp/modeling/layers/transformer_encoder_block.py,sha256=BiL8ErBs-m0UZ6ONVJV0ncfWX3LhMhPetIhfH2VvuD4,29910
+official/nlp/modeling/layers/transformer_encoder_block_test.py,sha256=g7oMDPvwg6Fv75SBdm6BInXPI8r5GcItBRjLFGuObyg,37821
 official/nlp/modeling/layers/transformer_scaffold.py,sha256=qmzhCJvbbFVF9zDqnfO4Zs2JDXwKhK7iEBOhsU6-KpQ,15704
 official/nlp/modeling/layers/transformer_scaffold_test.py,sha256=dRJwesTBKm-mF5mDHrHfVpVNnxa-Wx-fj_4ZHDPTpE0,19920
 official/nlp/modeling/layers/transformer_test.py,sha256=-pk9cdz9UlMpCIkGRkCKsMmjdRGi0seySaaB_2dwmXw,5522
@@ -1248,9 +1248,9 @@ tensorflow_models/tensorflow_models_test.py,sha256=yiAneltAW3NHSj3fUSvHNBjfq0MGZ
 tensorflow_models/nlp/__init__.py,sha256=8uQd4wI6Zc4IJMPjtQifMeWVbPFkTxqYh66wfivCOL4,807
 tensorflow_models/uplift/__init__.py,sha256=NzaweFf4ZmhRb2l_fuV6bP-2N8oSO3xu6xJqVb1UmpY,999
 tensorflow_models/vision/__init__.py,sha256=ks420Ooqzi0hU7HnQpM5rylLaE-YcJdJkBx_umVaXlE,833
-tf_models_nightly-2.20.0.dev20251028.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.20.0.dev20251028.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.20.0.dev20251028.dist-info/METADATA,sha256=uajPDCXLRHjjP9c1TuxbTh2QH1LBO84MLfydg3W5J50,1432
-tf_models_nightly-2.20.0.dev20251028.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.20.0.dev20251028.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.20.0.dev20251028.dist-info/RECORD,,
+tf_models_nightly-2.20.0.dev20251030.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.20.0.dev20251030.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.20.0.dev20251030.dist-info/METADATA,sha256=I2M0rXHA_2EU7FPchNNQCp9kEk7kXD5A35__qN9eB0E,1432
+tf_models_nightly-2.20.0.dev20251030.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.20.0.dev20251030.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.20.0.dev20251030.dist-info/RECORD,,