keras-hub-nightly 0.21.0.dev202505090407__py3-none-any.whl → 0.21.0.dev202505140407__py3-none-any.whl

keras_hub/src/layers/modeling/reversible_embedding.py

@@ -1,3 +1,5 @@
+import inspect
+
 import keras
 from keras import ops
 
@@ -184,31 +186,33 @@ class ReversibleEmbedding(keras.layers.Embedding):
         else:
             self._quantization_mode_error(self.quantization_mode)
 
-    def _int8_build(
-        self,
-        embeddings_initializer="zeros",
-        embeddings_scale_initializer="ones",
-        reverse_embeddings_initializer="zeros",
-        reverse_embeddings_scale_initializer="ones",
-    ):
-        super()._int8_build(
-            embeddings_initializer, embeddings_scale_initializer
-        )
+    def _int8_build(self, embeddings_shape=None):
+        if (
+            "embeddings_shape"
+            in inspect.signature(super()._int8_build).parameters
+        ):
+            if embeddings_shape is None:
+                embeddings_shape = (self.input_dim, self.output_dim)
+            super()._int8_build(embeddings_shape=embeddings_shape)
+        else:
+            # Backward compatibility for older versions of Keras.
+            super()._int8_build()
         self.inputs_quantizer = keras.quantizers.AbsMaxQuantizer(axis=-1)
         if not self.tie_weights:
             self.reverse_embeddings = self.add_weight(
                 name="reverse_embeddings",
                 shape=(self.output_dim, self.input_dim),
-                initializer=reverse_embeddings_initializer,
+                initializer="zeros",
                 dtype="int8",
                 trainable=False,
             )
             self.reverse_embeddings_scale = self.add_weight(
                 name="reverse_embeddings_scale",
                 shape=(self.input_dim,),
-                initializer=reverse_embeddings_scale_initializer,
+                initializer="ones",
                 trainable=False,
             )
+        self._is_quantized = True
 
     def _int8_call(self, inputs, reverse=False):
         if reverse:
@@ -232,27 +236,20 @@ class ReversibleEmbedding(keras.layers.Embedding):
         return super()._int8_call(inputs)
 
     def quantize(self, mode, type_check=True):
-        import gc
-
         if type_check and type(self) is not ReversibleEmbedding:
-            raise NotImplementedError(
-                f"Layer {self.__class__.__name__} does not have a `quantize()` "
-                "method implemented."
-            )
-        self._check_quantize_args(mode, self.compute_dtype)
+            raise self._not_implemented_error(self.quantize)
 
         def abs_max_quantize(inputs, axis):
             return keras.quantizers.abs_max_quantize(
                 inputs, axis=axis, to_numpy=True
             )
 
-        self._tracker.unlock()
+        embeddings_shape = (self.input_dim, self.output_dim)
         if mode == "int8":
             embeddings, embeddings_scale = abs_max_quantize(
                 self._embeddings, axis=-1
             )
             embeddings_scale = ops.squeeze(embeddings_scale, axis=-1)
-            self._untrack_variable(self._embeddings)
             del self._embeddings
             if not self.tie_weights:
                 reverse_embeddings, reverse_embeddings_scale = abs_max_quantize(
@@ -261,24 +258,17 @@ class ReversibleEmbedding(keras.layers.Embedding):
                 reverse_embeddings_scale = ops.squeeze(
                     reverse_embeddings_scale, axis=0
                 )
-                self._untrack_variable(self.reverse_embeddings)
                 del self.reverse_embeddings
-            else:
-                reverse_embeddings = None
-                reverse_embeddings_scale = None
-            self._int8_build(
-                lambda shape, dtype: embeddings,
-                lambda shape, dtype: embeddings_scale,
-                lambda shape, dtype: reverse_embeddings,
-                lambda shape, dtype: reverse_embeddings_scale,
-            )
-        else:
-            raise self._quantization_mode_error(mode)
-        self._tracker.lock()
+        self.quantized_build(embeddings_shape, mode)
+        if mode == "int8":
+            self._embeddings.assign(embeddings)
+            self.embeddings_scale.assign(embeddings_scale)
+            if not self.tie_weights:
+                self.reverse_embeddings.assign(reverse_embeddings)
+                self.reverse_embeddings_scale.assign(reverse_embeddings_scale)
 
         if self.dtype_policy.quantization_mode is None:
             policy = keras.dtype_policies.get(
                 f"{mode}_from_{self.dtype_policy.name}"
             )
             self.dtype_policy = policy
-        gc.collect()

keras_hub/src/version.py

@@ -1,7 +1,7 @@
 from keras_hub.src.api_export import keras_hub_export
 
 # Unique source of truth for the version number.
-__version__ = "0.21.0.dev202505090407"
+__version__ = "0.21.0.dev202505140407"
 
 
 @keras_hub_export("keras_hub.version")

{keras_hub_nightly-0.21.0.dev202505090407.dist-info → keras_hub_nightly-0.21.0.dev202505140407.dist-info}/METADATA

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: keras-hub-nightly
-Version: 0.21.0.dev202505090407
+Version: 0.21.0.dev202505140407
 Summary: Pretrained models for Keras.
 Author-email: Keras team <keras-users@googlegroups.com>
 License-Expression: Apache-2.0

{keras_hub_nightly-0.21.0.dev202505090407.dist-info → keras_hub_nightly-0.21.0.dev202505140407.dist-info}/RECORD

@@ -5,7 +5,7 @@ keras_hub/models/__init__.py,sha256=kFWNpjemQ8FLzDlFfMdAOOXJKtxuVHFxyZm7-1mH4Gc,
 keras_hub/samplers/__init__.py,sha256=aFQIkiqbZpi8vjrPp2MVII4QUfE-eQjra5fMeHsoy7k,886
 keras_hub/src/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_hub/src/api_export.py,sha256=9pQZK27JObxWZ96QPLBp1OBsjWigh1iuV6RglPGMRk0,1499
-keras_hub/src/version.py,sha256=0LAhJoL3V_nWN8KIbE6wSDHVx0bOL34Eu2MpW1o94Og,222
+keras_hub/src/version.py,sha256=byEOsjNvnRm-73Cv_UNklw4ZKk60II9uXyef5caMe0w,222
 keras_hub/src/layers/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_hub/src/layers/modeling/__init__.py,sha256=47DEQpj8HBSa-_TImW-5JCeuQeRkm5NMpJWZG3hSuFU,0
 keras_hub/src/layers/modeling/alibi_bias.py,sha256=1XBTHI52L_iJDhN_w5ydu_iMhCuTgQAxEPwcLA6BPuk,4411
@@ -16,7 +16,7 @@ keras_hub/src/layers/modeling/f_net_encoder.py,sha256=zkVeO5Nk_kBZCUGq2LeDGmPEIM
 keras_hub/src/layers/modeling/masked_lm_head.py,sha256=no6XQb76KB2cUiksYC0MSdyeDOK7pn8MY6cmdCDxpKs,9015
 keras_hub/src/layers/modeling/non_max_supression.py,sha256=yAkAH1CCj_tYXgQTav39IRr_Uxn8hmzJgIxqbYQyZY8,22565
 keras_hub/src/layers/modeling/position_embedding.py,sha256=FfTS6JGMhnOIzo9bHzvoxBbdQNctc32iRLI7ZjdxoTY,3850
-keras_hub/src/layers/modeling/reversible_embedding.py,sha256=sfm5giI-bHu2J9xm9Tkydx8XM-I_m8Oe0wbW1gzrYjk,11141
+keras_hub/src/layers/modeling/reversible_embedding.py,sha256=w6f1LQzwPOKUdlWDy3YRECaDzb8veCB2PAxy4L0HJ7w,10866
 keras_hub/src/layers/modeling/rms_normalization.py,sha256=Ylnc9vkDw1A_ZqiKpQ09jVTAGumS5rspjdQOkH-mxf4,1084
 keras_hub/src/layers/modeling/rotary_embedding.py,sha256=BuMD2dCyZi73Eokddx8Q9cFb4pJVlOL2OgFwsom2p8I,6059
 keras_hub/src/layers/modeling/sine_position_encoding.py,sha256=NAPW9HaVTMNZgUJNzA3l1B3C_FNvaY7IW-5tQgFgnNg,3453
@@ -486,7 +486,7 @@ keras_hub/src/utils/transformers/preset_loader.py,sha256=1nfS5xVsl-JROGXJXltTqV1
 keras_hub/src/utils/transformers/safetensor_utils.py,sha256=CYUHyA4y-B61r7NDnCsFb4t_UmSwZ1k9L-8gzEd6KRg,3339
 keras_hub/tokenizers/__init__.py,sha256=4etC--bzhczJrRcvCmxZmOC9hJJcIVOUCgmqMLB3bp0,4051
 keras_hub/utils/__init__.py,sha256=jXPqVGBpJr_PpYmqD8aDG-fRMlxH-ulqCR2SZMn288Y,646
-keras_hub_nightly-0.21.0.dev202505090407.dist-info/METADATA,sha256=zjzu6HMAt2i4YdX8hB6AnbFxGMSTHg2fClV2KACeBYs,7393
-keras_hub_nightly-0.21.0.dev202505090407.dist-info/WHEEL,sha256=0CuiUZ_p9E4cD6NyLD6UG80LBXYyiSYZOKDm5lp32xk,91
-keras_hub_nightly-0.21.0.dev202505090407.dist-info/top_level.txt,sha256=N4J6piIWBKa38A4uV-CnIopnOEf8mHAbkNXafXm_CuA,10
-keras_hub_nightly-0.21.0.dev202505090407.dist-info/RECORD,,
+keras_hub_nightly-0.21.0.dev202505140407.dist-info/METADATA,sha256=uD895R3mp5irmKsF6VsQuOJQ92uDjSDrBelTy93WHx8,7393
+keras_hub_nightly-0.21.0.dev202505140407.dist-info/WHEEL,sha256=DnLRTWE75wApRYVsjgc6wsVswC54sMSJhAEd4xhDpBk,91
+keras_hub_nightly-0.21.0.dev202505140407.dist-info/top_level.txt,sha256=N4J6piIWBKa38A4uV-CnIopnOEf8mHAbkNXafXm_CuA,10
+keras_hub_nightly-0.21.0.dev202505140407.dist-info/RECORD,,

{keras_hub_nightly-0.21.0.dev202505090407.dist-info → keras_hub_nightly-0.21.0.dev202505140407.dist-info}/WHEEL

@@ -1,5 +1,5 @@
 Wheel-Version: 1.0
-Generator: setuptools (80.3.1)
+Generator: setuptools (80.4.0)
 Root-Is-Purelib: true
 Tag: py3-none-any