tf-models-nightly 2.19.0.dev20250102__py2.py3-none-any.whl → 2.19.0.dev20250103__py2.py3-none-any.whl

official/nlp/modeling/layers/block_sparse_attention.py

@@ -84,6 +84,17 @@ class MultiHeadAttention(tf_keras.layers.MultiHeadAttention):
             "sigmoid_attn_bias must be specified for sigmoid attn."
         )
 
+  def get_config(self):
+    config = super().get_config()
+    config.update({
+        "src_block_size": self._src_block_size,
+        "tgt_block_size": self._tgt_block_size,
+        "use_sigmoid_attn": self._use_sigmoid_attn,
+        "sigmoid_attn_bias": self._sigmoid_attn_bias,
+        "num_kv_heads": self._num_kv_heads,
+    })
+    return config
+
   def _build_from_signature(self, query, value, key=None):
     # pytype: disable=attribute-error
     super()._build_from_signature(query, value, key)

official/nlp/modeling/layers/multi_query_attention.py

@@ -93,6 +93,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],

{tf_models_nightly-2.19.0.dev20250102.dist-info → tf_models_nightly-2.19.0.dev20250103.dist-info}/METADATA

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

{tf_models_nightly-2.19.0.dev20250102.dist-info → tf_models_nightly-2.19.0.dev20250103.dist-info}/RECORD

@@ -305,7 +305,7 @@ official/nlp/modeling/layers/bigbird_attention.py,sha256=dzutgRoQt2DFsYMpMILv_QF
 official/nlp/modeling/layers/bigbird_attention_test.py,sha256=cBYwK5k1rnykZ0gif-n7VaByLIoElA-N0_svCRKASoU,2206
 official/nlp/modeling/layers/block_diag_feedforward.py,sha256=FDEt-J_QjOxwar3eT5yjMs4hR41Ppke1zj7iswsZR4M,7243
 official/nlp/modeling/layers/block_diag_feedforward_test.py,sha256=wcg8In6FIOCxcKqe5rucftjJ_kUWTi9Ei7eEmlVCYpE,4181
-official/nlp/modeling/layers/block_sparse_attention.py,sha256=8Jyinyf5XuU6nuxblGRfNnAviBEZIltMSoNJzAVjAYo,12233
+official/nlp/modeling/layers/block_sparse_attention.py,sha256=3UDjvvv4oYM-0oLUu_jJNueKfwiFm4sAT-qymicG3kU,12583
 official/nlp/modeling/layers/block_sparse_attention_test.py,sha256=9YiKtv4YCrKIyUbv27P2xcTXFohoaRxq2K6vsOUi4zU,16447
 official/nlp/modeling/layers/cls_head.py,sha256=0X_gdjnAt6TZVrH_xkDcQCpwLuVz5Pb7d04wEVN_Kn8,16208
 official/nlp/modeling/layers/cls_head_test.py,sha256=01oMmiuyp1lDEXBYa9r3krn6BtH-QuSedGOca9LViEc,8888
@@ -331,7 +331,7 @@ official/nlp/modeling/layers/moe.py,sha256=azoK1Cz5l-5yT-FOtb7VYSbePTxmoje-2XAcq
 official/nlp/modeling/layers/moe_test.py,sha256=sr4hys2su_fcf9-6t0awqOkico92pQ3I2ApVF2rvvf0,9414
 official/nlp/modeling/layers/multi_channel_attention.py,sha256=fInxZUAPrrgCoCkM62JylM_vf5Wp8D0uZAvd5yuooVU,7322
 official/nlp/modeling/layers/multi_channel_attention_test.py,sha256=KcKuq22gOOq8HBzxRs3gGCwDRhmJPzYxfzfnqorQ1fw,1922
-official/nlp/modeling/layers/multi_query_attention.py,sha256=fFPBa9IBVj_O5x5OfGuHUFnJmiouNL6F1KsBCeHUqwM,6978
+official/nlp/modeling/layers/multi_query_attention.py,sha256=efi58Dne1kDSH6TdfTHW9CJVVOJB10fSvvQhC-8JUHE,7111
 official/nlp/modeling/layers/multi_query_attention_test.py,sha256=3VFF2hz85YExWPwdbhYWaSrIaSOkC1x7axdGfXr0W90,8512
 official/nlp/modeling/layers/on_device_embedding.py,sha256=FgsHyRXf5TWVTyo4OeKImmrTnn4uOPJgS3AGKzKMWYY,4582
 official/nlp/modeling/layers/on_device_embedding_test.py,sha256=M-LUba4QXV37s9Cx7aH8LL3bz_YotC6qITmWRI7Fhjk,8589
@@ -1222,9 +1222,9 @@ tensorflow_models/tensorflow_models_test.py,sha256=nc6A9K53OGqF25xN5St8EiWvdVbda
 tensorflow_models/nlp/__init__.py,sha256=4tA5Pf4qaFwT-fIFOpX7x7FHJpnyJT-5UgOeFYTyMlc,807
 tensorflow_models/uplift/__init__.py,sha256=mqfa55gweOdpKoaQyid4A_4u7xw__FcQeSIF0k_pYmI,999
 tensorflow_models/vision/__init__.py,sha256=zBorY_v5xva1uI-qxhZO3Qh-Dii-Suq6wEYh6hKHDfc,833
-tf_models_nightly-2.19.0.dev20250102.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
-tf_models_nightly-2.19.0.dev20250102.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
-tf_models_nightly-2.19.0.dev20250102.dist-info/METADATA,sha256=wpz6mIqvedDtyUHDM4pKQfHos2FIhmOE93NmfQvIyiY,1432
-tf_models_nightly-2.19.0.dev20250102.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
-tf_models_nightly-2.19.0.dev20250102.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
-tf_models_nightly-2.19.0.dev20250102.dist-info/RECORD,,
+tf_models_nightly-2.19.0.dev20250103.dist-info/AUTHORS,sha256=1dG3fXVu9jlo7bul8xuix5F5vOnczMk7_yWn4y70uw0,337
+tf_models_nightly-2.19.0.dev20250103.dist-info/LICENSE,sha256=WxeBS_DejPZQabxtfMOM_xn8qoZNJDQjrT7z2wG1I4U,11512
+tf_models_nightly-2.19.0.dev20250103.dist-info/METADATA,sha256=OzGqZyxnQDnpq6JstzJEuJPHLq31ivjdzj-gLfYzi04,1432
+tf_models_nightly-2.19.0.dev20250103.dist-info/WHEEL,sha256=kGT74LWyRUZrL4VgLh6_g12IeVl_9u9ZVhadrgXZUEY,110
+tf_models_nightly-2.19.0.dev20250103.dist-info/top_level.txt,sha256=gum2FfO5R4cvjl2-QtP-S1aNmsvIZaFFT6VFzU0f4-g,33
+tf_models_nightly-2.19.0.dev20250103.dist-info/RECORD,,