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

official/nlp/modeling/layers/rezero_transformer.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.
@@ -21,6 +21,8 @@ import gin
 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 util
 
 
@@ -53,6 +55,12 @@ class ReZeroTransformer(tf_keras.layers.Layer):
     bias_constraint: Constraint for dense layer kernels.
     use_layer_norm: If add layer_norm on top of the ReZero.
     share_rezero: If attention layer and FFN layer share the same alpha.
+    num_kv_heads: Number of key-value heads for multi-query attention. Refer to
+      `multi_query_attention.MultiHeadAttention` for more details.
+    src_block_size: Source block size. Refer to
+      `block_sparse_attention.MultiHeadAttention` for more details.
+    tgt_block_size: Target block size. Refer to
+      `block_sparse_attention.MultiHeadAttention` for more details.
   """
 
   def __init__(self,
@@ -71,6 +79,13 @@ class ReZeroTransformer(tf_keras.layers.Layer):
                bias_constraint=None,
                use_layer_norm=False,
                share_rezero=True,
+               num_kv_heads=None,
+               src_block_size=None,
+               tgt_block_size=None,
+               linformer_dim=None,
+               linformer_shared_kv_projection=True,
+               use_sigmoid_attn=False,
+               sigmoid_attn_bias=None,
                **kwargs):
     # attention_dropout will override attention_dropout_rate.
     # This is to unify the input params with TransformerEncoderBlock.
@@ -101,6 +116,23 @@ class ReZeroTransformer(tf_keras.layers.Layer):
     self._bias_constraint = tf_keras.constraints.get(bias_constraint)
     self._use_layer_norm = use_layer_norm
     self._share_rezero = share_rezero
+    self._num_kv_heads = num_kv_heads
+    self._src_block_size = src_block_size
+    self._tgt_block_size = tgt_block_size
+    self._linformer_dim = linformer_dim
+    self._linformer_shared_kv_projection = linformer_shared_kv_projection
+    self._use_sigmoid_attn = use_sigmoid_attn
+    self._sigmoid_attn_bias = sigmoid_attn_bias
+    if self._linformer_dim is not None or self._use_sigmoid_attn:
+      raise ValueError(
+          "Linformer and Sigmoid attention are not supported in ReZero"
+          " Transformer."
+      )
+    if self._num_kv_heads is not None and self._src_block_size is not None:
+      raise ValueError(
+          "Block sparse attention does not support Multi-query attention."
+          " Specify only one of them."
+      )
 
   def build(self, input_shape):
     if isinstance(input_shape, tf.TensorShape):
@@ -109,53 +141,77 @@ class ReZeroTransformer(tf_keras.layers.Layer):
       input_tensor_shape = tf.TensorShape(input_shape[0])
     else:
       raise ValueError(
-          "The type of input shape argument is not supported, got: %s" %
-          type(input_shape))
+          "The type of input shape argument is not supported, got: %s"
+          % type(input_shape)
+      )
 
     if len(input_tensor_shape.as_list()) != 3:
-      raise ValueError("TransformerLayer expects a three-dimensional input of "
-                       "shape [batch, sequence, width].")
+      raise ValueError(
+          "TransformerLayer expects a three-dimensional input of "
+          "shape [batch, sequence, width]."
+      )
     batch_size, sequence_length, hidden_size = input_tensor_shape
 
     if len(input_shape) == 2:
       mask_tensor_shape = tf.TensorShape(input_shape[1])
       expected_mask_tensor_shape = tf.TensorShape(
-          [batch_size, sequence_length, sequence_length])
+          [batch_size, sequence_length, sequence_length]
+      )
       if not expected_mask_tensor_shape.is_compatible_with(mask_tensor_shape):
-        raise ValueError("When passing a mask tensor to TransformerLayer, the "
-                         "mask tensor must be of shape [batch, "
-                         "sequence_length, sequence_length] (here %s). Got a "
-                         "mask tensor of shape %s." %
-                         (expected_mask_tensor_shape, mask_tensor_shape))
+        raise ValueError(
+            "When passing a mask tensor to TransformerLayer, the "
+            "mask tensor must be of shape [batch, "
+            "sequence_length, sequence_length] (here %s). Got a "
+            "mask tensor of shape %s."
+            % (expected_mask_tensor_shape, mask_tensor_shape)
+        )
     if hidden_size % self._num_heads != 0:
       raise ValueError(
           "The input size (%d) is not a multiple of the number of attention "
-          "heads (%d)" % (hidden_size, self._num_heads))
+          "heads (%d)" % (hidden_size, self._num_heads)
+      )
     self._attention_head_size = int(hidden_size // self._num_heads)
     common_kwargs = dict(
         kernel_regularizer=self._kernel_regularizer,
         bias_regularizer=self._bias_regularizer,
         activity_regularizer=self._activity_regularizer,
         kernel_constraint=self._kernel_constraint,
-        bias_constraint=self._bias_constraint)
-    self._attention_layer = tf_keras.layers.MultiHeadAttention(
+        bias_constraint=self._bias_constraint,
+    )
+    attention_kwargs = dict(
         num_heads=self._num_heads,
         key_dim=self._attention_head_size,
         dropout=self._attention_dropout_rate,
         name="self_attention",
         kernel_initializer=tf_utils.clone_initializer(self._kernel_initializer),
         bias_initializer=tf_utils.clone_initializer(self._bias_initializer),
-        **common_kwargs)
+    )
+    if self._src_block_size is not None:
+      attention_kwargs.update(
+          src_block_size=self._src_block_size,
+          tgt_block_size=self._tgt_block_size,
+          name="block_sparse_attention",
+      )
+      attention_fn = block_sparse_attention.MultiHeadAttention
+    elif self._num_kv_heads is not None:
+      attention_kwargs.update(
+          num_kv_heads=self._num_kv_heads,
+          name="multi_query_attention",
+          )
+      attention_fn = multi_query_attention.MultiHeadAttention
+    else:
+      attention_fn = tf_keras.layers.MultiHeadAttention
+    self._attention_layer = attention_fn(**attention_kwargs, **common_kwargs)
     self._attention_dropout = tf_keras.layers.Dropout(rate=self._dropout_rate)
     if self._use_layer_norm:
       # Use float32 in layernorm for numeric stability.
       # It is probably safe in mixed_float16, but we haven't validated this yet.
-      self._attention_layer_norm = (
-          tf_keras.layers.LayerNormalization(
-              name="self_attention_layer_norm",
-              axis=-1,
-              epsilon=1e-12,
-              dtype=tf.float32))
+      self._attention_layer_norm = tf_keras.layers.LayerNormalization(
+          name="self_attention_layer_norm",
+          axis=-1,
+          epsilon=1e-12,
+          dtype=tf.float32,
+      )
     self._intermediate_dense = tf_keras.layers.EinsumDense(
         "abc,cd->abd",
         output_shape=(None, self._inner_dim),
@@ -221,6 +277,12 @@ class ReZeroTransformer(tf_keras.layers.Layer):
             self._use_layer_norm,
         "share_rezero":
             self._share_rezero,
+        "num_kv_heads":
+            self._num_kv_heads,
+        "src_block_size":
+            self._src_block_size,
+        "tgt_block_size":
+            self._tgt_block_size,
         "kernel_initializer":
             tf_keras.initializers.serialize(self._kernel_initializer),
         "bias_initializer":
@@ -235,6 +297,12 @@ class ReZeroTransformer(tf_keras.layers.Layer):
             tf_keras.constraints.serialize(self._kernel_constraint),
         "bias_constraint":
             tf_keras.constraints.serialize(self._bias_constraint),
+        "linformer_dim": self._linformer_dim,
+        "linformer_shared_kv_projection": (
+            self._linformer_shared_kv_projection
+        ),
+        "use_sigmoid_attn": self._use_sigmoid_attn,
+        "sigmoid_attn_bias": self._sigmoid_attn_bias,
     }
     base_config = super().get_config()
     return dict(list(base_config.items()) + list(config.items()))

official/nlp/modeling/layers/rezero_transformer_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.
@@ -141,6 +141,69 @@ class TransformerWithReZeroLayerTest(tf.test.TestCase, parameterized.TestCase):
     output = test_layer(inputs)
     self.assertEqual(output.shape, q_tensor.shape)
 
+  @parameterized.named_parameters(('_mqa', 1),
+                                  ('_gqa', 5))
+  def test_rezero_with_kv_heads(self, num_kv_heads):
+    tf_keras.mixed_precision.set_global_policy('mixed_float16')
+    test_layer = rezero_transformer.ReZeroTransformer(
+        num_attention_heads=10,
+        intermediate_size=2048,
+        intermediate_activation='relu',
+        num_kv_heads=num_kv_heads,
+        )
+    sequence_length = 21
+    width = 80
+    # Create a 3-dimensional input (the first dimension is implicit).
+    data_tensor = tf_keras.Input(shape=(sequence_length, width))
+    # Create a 2-dimensional input (the first dimension is implicit).
+    mask_tensor = tf_keras.Input(shape=(sequence_length, sequence_length))
+    output_tensor = test_layer([data_tensor, mask_tensor])
+
+    # Create a model from the test layer.
+    model = tf_keras.Model([data_tensor, mask_tensor], output_tensor)
+
+    # Invoke the model on test data. We can't validate the output data itself
+    # (the NN is too complex) but this will rule out structural runtime errors.
+    batch_size = 6
+    input_data = (10 * np.random.random_sample(
+        (batch_size, sequence_length, width)))
+    # The attention mask should be of shape (batch, from_seq_len, to_seq_len),
+    # which here is (batch, sequence_length, sequence_length)
+    mask_data = np.random.randint(
+        2, size=(batch_size, sequence_length, sequence_length))
+    _ = model.predict([input_data, mask_data])
+
+  def test_rezero_with_block_sparse_attention(self):
+    tf_keras.mixed_precision.set_global_policy('mixed_float16')
+    test_layer = rezero_transformer.ReZeroTransformer(
+        num_attention_heads=10,
+        intermediate_size=2048,
+        intermediate_activation='relu',
+        src_block_size=3,
+        tgt_block_size=3,
+        )
+    sequence_length = 21
+    width = 80
+    # Create a 3-dimensional input (the first dimension is implicit).
+    data_tensor = tf_keras.Input(shape=(sequence_length, width))
+    # Create a 2-dimensional input (the first dimension is implicit).
+    mask_tensor = tf_keras.Input(shape=(sequence_length, sequence_length))
+    output_tensor = test_layer([data_tensor, mask_tensor])
+
+    # Create a model from the test layer.
+    model = tf_keras.Model([data_tensor, mask_tensor], output_tensor)
+
+    # Invoke the model on test data. We can't validate the output data itself
+    # (the NN is too complex) but this will rule out structural runtime errors.
+    batch_size = 6
+    input_data = (10 * np.random.random_sample(
+        (batch_size, sequence_length, width)))
+    # The attention mask should be of shape (batch, from_seq_len, to_seq_len),
+    # which here is (batch, sequence_length, sequence_length)
+    mask_data = np.random.randint(
+        2, size=(batch_size, sequence_length, sequence_length))
+    _ = model.predict([input_data, mask_data])
+
 
 if __name__ == '__main__':
   tf.test.main()

official/nlp/modeling/layers/routing.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/routing_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/self_attention_mask.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/spectral_normalization.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/spectral_normalization_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/talking_heads_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/talking_heads_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/text_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/text_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/tn_expand_condense.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/tn_expand_condense_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/tn_transformer_expand_condense.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,8 +14,6 @@
 
 """TN-BERT TNTransformerExpandCondense employing Expand-Condense layer instead of Dense."""
 # pylint: disable=g-classes-have-attributes
-# Import libraries
-
 import gin
 import tensorflow as tf, tf_keras
 

official/nlp/modeling/layers/tn_transformer_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/transformer.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.