tf-models-nightly 2.20.0.dev20251114__py2.py3-none-any.whl → 2.20.0.dev20251116__py2.py3-none-any.whl

official/nlp/modeling/layers/transformer_encoder_block.py

@@ -121,6 +121,7 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
       lowrank_query_seq_proj_dim=None,
       enable_talking_heads=False,
       enable_gqa_optimization=False,
+      softmax_robust_masking=False,
       **kwargs,
   ):
     """Initializes `TransformerEncoderBlock`.
@@ -209,6 +210,8 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
         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.
+      softmax_robust_masking: If true, will use a more numerically robust
+        masking impl for softmax.
       **kwargs: keyword arguments.
     """
     util.filter_kwargs(kwargs)
@@ -253,6 +256,7 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
     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
+    self._softmax_robust_masking = softmax_robust_masking
     if (
         self._src_block_size is not None
         and self._num_kv_heads is not None
@@ -314,6 +318,7 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
         bias_initializer=tf_utils.clone_initializer(self._bias_initializer),
         attention_axes=self._attention_axes,
         output_shape=self._output_last_dim,
+        softmax_robust_masking=self._softmax_robust_masking,
         name="self_attention",
     )
     common_kwargs = dict(
@@ -512,6 +517,7 @@ class TransformerEncoderBlock(tf_keras.layers.Layer):
         "linformer_dim": self._linformer_dim,
         "linformer_shared_kv_projection": self._linformer_shared_kv_projection,
         "lowrank_query_seq_proj_dim": self._lowrank_query_seq_proj_dim,
+        "softmax_robust_masking": self._softmax_robust_masking
     }
     base_config = super().get_config()
     return dict(list(base_config.items()) + list(config.items()))

{tf_models_nightly-2.20.0.dev20251114.dist-info → tf_models_nightly-2.20.0.dev20251116.dist-info}/METADATA

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

{tf_models_nightly-2.20.0.dev20251114.dist-info → tf_models_nightly-2.20.0.dev20251116.dist-info}/RECORD

@@ -363,7 +363,7 @@ 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=BiL8ErBs-m0UZ6ONVJV0ncfWX3LhMhPetIhfH2VvuD4,29910
+official/nlp/modeling/layers/transformer_encoder_block.py,sha256=78nIg3hVFKLeqJ-Hdg4j1scfskZAnm2CSHqlTeoSRcc,30236
 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
@@ -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.dev20251114.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.20.0.dev20251114.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.20.0.dev20251114.dist-info/METADATA,sha256=0ImfIyWGCsTpCzpIt8ishPQk_wsbj7V3CgYr99XgD1A,1432
-tf_models_nightly-2.20.0.dev20251114.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.20.0.dev20251114.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.20.0.dev20251114.dist-info/RECORD,,
+tf_models_nightly-2.20.0.dev20251116.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.20.0.dev20251116.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.20.0.dev20251116.dist-info/METADATA,sha256=O3KI3rsF_AQuFE494hYOCSC1ZiYV9Bf8BZ09QYXO2m4,1432
+tf_models_nightly-2.20.0.dev20251116.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.20.0.dev20251116.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.20.0.dev20251116.dist-info/RECORD,,