tf-models-nightly 2.20.0.dev20251029__py2.py3-none-any.whl → 2.20.0.dev20251031__py2.py3-none-any.whl

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.dev20251029.dist-info → tf_models_nightly-2.20.0.dev20251031.dist-info}/METADATA

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

{tf_models_nightly-2.20.0.dev20251029.dist-info → tf_models_nightly-2.20.0.dev20251031.dist-info}/RECORD

@@ -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.dev20251029.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.20.0.dev20251029.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.20.0.dev20251029.dist-info/METADATA,sha256=hTl-klu4MFPyEUTW4FX3TyByeXJs4Hp7BJpSjTWFaCw,1432
-tf_models_nightly-2.20.0.dev20251029.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.20.0.dev20251029.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.20.0.dev20251029.dist-info/RECORD,,
+tf_models_nightly-2.20.0.dev20251031.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.20.0.dev20251031.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.20.0.dev20251031.dist-info/METADATA,sha256=Nkc6PAbFGnKBfKCfhkqghJ4risJv55iE9dKZblkfEq0,1432
+tf_models_nightly-2.20.0.dev20251031.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.20.0.dev20251031.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.20.0.dev20251031.dist-info/RECORD,,