Try requiring a miniumum version of keras

.github/workflows/actions.yml

@@ -20,7 +20,7 @@ jobs:
         version: [keras-stable]
         include:
           - backend: jax
-            version: keras-3.1
+            version: keras-3.5
           - backend: jax
             version: keras-nightly
     runs-on: ubuntu-latest
@@ -48,11 +48,11 @@ jobs:
       run: |
           pip install -r requirements.txt --progress-bar off
           pip install --no-deps -e "." --progress-bar off
-    - name: Pin Keras 3.1
-      if: ${{ matrix.version == 'keras-3.1'}}
+    - name: Pin Keras 3.5
+      if: ${{ matrix.version == 'keras-3.5'}}
       run: |
         pip uninstall -y keras
-        pip install keras==3.1.0 --progress-bar off
+        pip install keras==3.5.0 --progress-bar off
     - name: Pin Keras Nightly
       if: ${{ matrix.version == 'keras-nightly'}}
       run: |

keras_hub/src/layers/modeling/reversible_embedding.py

@@ -1,9 +1,7 @@
 import keras
 from keras import ops
-from packaging.version import parse
 
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.utils.keras_utils import assert_quantization_support
 
 
 @keras_hub_export("keras_hub.layers.ReversibleEmbedding")
@@ -145,10 +143,6 @@ def save_own_variables(self, store):
         if not self.built:
             return
         super().save_own_variables(store)
-        # Before Keras 3.2, the reverse weight is saved in the super() call.
-        # After Keras 3.2, the reverse weight must be saved manually.
-        if parse(keras.version()) < parse("3.2.0"):
-            return
         target_variables = []
         if not self.tie_weights:
             # Store the reverse embedding weights as the last weights.
@@ -239,9 +233,7 @@ def _int8_call(self, inputs, reverse=False):
 
     def quantize(self, mode, type_check=True):
         import gc
-        import inspect
 
-        assert_quantization_support()
         if type_check and type(self) is not ReversibleEmbedding:
             raise NotImplementedError(
                 f"Layer {self.__class__.__name__} does not have a `quantize()` "
@@ -250,14 +242,9 @@ def quantize(self, mode, type_check=True):
         self._check_quantize_args(mode, self.compute_dtype)
 
         def abs_max_quantize(inputs, axis):
-            sig = inspect.signature(keras.quantizers.abs_max_quantize)
-            if "to_numpy" in sig.parameters:
-                return keras.quantizers.abs_max_quantize(
-                    inputs, axis=axis, to_numpy=True
-                )
-            else:
-                # `keras<=3.4.1` doesn't support `to_numpy`
-                return keras.quantizers.abs_max_quantize(inputs, axis=axis)
+            return keras.quantizers.abs_max_quantize(
+                inputs, axis=axis, to_numpy=True
+            )
 
         self._tracker.unlock()
         if mode == "int8":

keras_hub/src/layers/modeling/reversible_embedding_test.py

@@ -10,7 +10,6 @@
     ReversibleEmbedding,
 )
 from keras_hub.src.tests.test_case import TestCase
-from keras_hub.src.utils.keras_utils import has_quantization_support
 
 
 class ReversibleEmbeddingTest(TestCase):
@@ -97,9 +96,6 @@ def test_reverse_dtype(self):
         ("tie_weights", True), ("untie_weights", False)
     )
     def test_quantize_int8(self, tie_weights):
-        if not has_quantization_support():
-            self.skipTest("This version of Keras doesn't support quantization.")
-
         layer_config = dict(
             input_dim=100, output_dim=32, tie_weights=tie_weights
         )
@@ -148,9 +144,6 @@ def test_quantize_int8(self, tie_weights):
         ("untie_weights", False),
     )
     def test_quantize_dtype_argument(self, tie_weights):
-        if not has_quantization_support():
-            self.skipTest("This version of Keras doesn't support quantization.")
-
         self.run_layer_test(
             cls=ReversibleEmbedding,
             init_kwargs={

keras_hub/src/models/backbone.py

@@ -1,7 +1,6 @@
 import keras
 
 from keras_hub.src.api_export import keras_hub_export
-from keras_hub.src.utils.keras_utils import assert_quantization_support
 from keras_hub.src.utils.preset_utils import builtin_presets
 from keras_hub.src.utils.preset_utils import get_preset_loader
 from keras_hub.src.utils.preset_utils import get_preset_saver
@@ -83,10 +82,6 @@ def token_embedding(self):
     def token_embedding(self, value):
         self._token_embedding = value
 
-    def quantize(self, mode, **kwargs):
-        assert_quantization_support()
-        return super().quantize(mode, **kwargs)
-
     def get_config(self):
         # Don't chain to super here. `get_config()` for functional models is
         # a nested layer config and cannot be passed to Backbone constructors.

keras_hub/src/tests/test_case.py

@@ -15,7 +15,6 @@
 )
 from keras_hub.src.models.retinanet.feature_pyramid import FeaturePyramid
 from keras_hub.src.tokenizers.tokenizer import Tokenizer
-from keras_hub.src.utils.keras_utils import has_quantization_support
 from keras_hub.src.utils.tensor_utils import is_float_dtype
 
 
@@ -487,7 +486,7 @@ def run_backbone_test(
             self.run_precision_test(cls, init_kwargs, input_data)
 
         # Check quantization.
-        if run_quantization_check and has_quantization_support():
+        if run_quantization_check:
             self.run_quantization_test(backbone, cls, init_kwargs, input_data)
 
     def run_vision_backbone_test(

keras_hub/src/utils/keras_utils.py

@@ -2,7 +2,6 @@
 
 import keras
 from absl import logging
-from packaging.version import parse
 
 try:
     import tensorflow as tf
@@ -41,18 +40,6 @@ def gelu_approximate(x):
     return keras.activations.gelu(x, approximate=True)
 
 
-def has_quantization_support():
-    return False if parse(keras.version()) < parse("3.4.0") else True
-
-
-def assert_quantization_support():
-    if not has_quantization_support():
-        raise ValueError(
-            "Quantization API requires Keras >= 3.4.0 to function "
-            f"correctly. Received: '{keras.version()}'"
-        )
-
-
 def standardize_data_format(data_format):
     if data_format is None:
         return keras.config.image_data_format()

keras_hub/src/utils/preset_utils.py

@@ -7,20 +7,10 @@
 
 import keras
 from absl import logging
-from packaging.version import parse
 
 from keras_hub.src.api_export import keras_hub_export
 from keras_hub.src.utils.keras_utils import print_msg
 
-try:
-    import tensorflow as tf
-except ImportError:
-    raise ImportError(
-        "To use `keras_hub`, please install Tensorflow: "
-        "`pip install tensorflow`. The TensorFlow package is required for data "
-        "preprocessing with any backend."
-    )
-
 try:
     import kagglehub
     from kagglehub.exceptions import KaggleApiHTTPError
@@ -173,22 +163,8 @@ def get_file(preset, path):
                 )
             else:
                 raise ValueError(message)
-
-    elif scheme in tf.io.gfile.get_registered_schemes():
-        url = os.path.join(preset, path)
-        subdir = preset.replace("://", "_").replace("-", "_").replace("/", "_")
-        filename = os.path.basename(path)
-        subdir = os.path.join(subdir, os.path.dirname(path))
-        try:
-            return copy_gfile_to_cache(
-                filename,
-                url,
-                cache_subdir=os.path.join("models", subdir),
-            )
-        except (tf.errors.PermissionDeniedError, tf.errors.NotFoundError) as e:
-            raise FileNotFoundError(
-                f"`{path}` doesn't exist in preset directory `{preset}`.",
-            ) from e
+    elif scheme in tf_registered_schemes():
+        return tf_copy_gfile_to_cache(preset, path)
     elif scheme == HF_SCHEME:
         if huggingface_hub is None:
             raise ImportError(
@@ -237,26 +213,48 @@ def get_file(preset, path):
         )
 
 
-def copy_gfile_to_cache(filename, url, cache_subdir):
+def tf_registered_schemes():
+    try:
+        import tensorflow as tf
+
+        return tf.io.gfile.get_registered_schemes()
+    except ImportError:
+        return []
+
+
+def tf_copy_gfile_to_cache(preset, path):
     """Much of this is adapted from get_file of keras core."""
     if "KERAS_HOME" in os.environ:
-        cachdir_base = os.environ.get("KERAS_HOME")
+        basedir = os.environ.get("KERAS_HOME")
     else:
-        cachdir_base = os.path.expanduser(os.path.join("~", ".keras"))
-    if not os.access(cachdir_base, os.W_OK):
-        cachdir_base = os.path.join("/tmp", ".keras")
-    cachedir = os.path.join(cachdir_base, cache_subdir)
+        basedir = os.path.expanduser(os.path.join("~", ".keras"))
+    if not os.access(basedir, os.W_OK):
+        basedir = os.path.join("/tmp", ".keras")
+
+    url = os.path.join(preset, path)
+    subdir = preset.replace("://", "_").replace("-", "_").replace("/", "_")
+    filename = os.path.basename(path)
+    subdir = os.path.join(subdir, os.path.dirname(path))
+    cachedir = os.path.join(basedir, "models", subdir, os.path.dirname(path))
     os.makedirs(cachedir, exist_ok=True)
 
     fpath = os.path.join(cachedir, filename)
     if not os.path.exists(fpath):
         print_msg(f"Downloading data from {url}")
         try:
+            import tensorflow as tf
+
             tf.io.gfile.copy(url, fpath)
         except Exception as e:
             # gfile.copy will leave an empty file after an error.
             # Work around this bug.
             os.remove(fpath)
+            if isinstance(
+                e, tf.errors.PermissionDeniedError, tf.errors.NotFoundError
+            ):
+                raise FileNotFoundError(
+                    f"`{path}` doesn't exist in preset directory `{preset}`.",
+                ) from e
             raise e
 
     return fpath
@@ -394,12 +392,6 @@ def upload_preset(
                 "Uploading a model to Kaggle Hub requires the `kagglehub` "
                 "package. Please install with `pip install kagglehub`."
             )
-        if parse(kagglehub.__version__) < parse("0.2.4"):
-            raise ImportError(
-                "Uploading a model to Kaggle Hub requires the `kagglehub` "
-                "package version `0.2.4` or higher. Please upgrade with "
-                "`pip install --upgrade kagglehub`."
-            )
         kaggle_handle = uri.removeprefix(KAGGLE_PREFIX)
         kagglehub.model_upload(kaggle_handle, preset)
     elif uri.startswith(HF_PREFIX):

setup.py

@@ -38,17 +38,14 @@ def get_version(rel_path):
     author_email="[email protected]",
     license="Apache License 2.0",
     install_requires=[
+        "keras>=3.4.0",
         "absl-py",
         "numpy",
         "packaging",
         "regex",
         "rich",
         "kagglehub",
-        # Don't require tensorflow-text on MacOS, there are no binaries for ARM.
-        # Also, we rely on tensorflow *transitively* through tensorflow-text.
-        # This avoid a slowdown during `pip install keras-hub` where pip would
-        # download many version of both libraries to find compatible versions.
-        "tensorflow-text; platform_system != 'Darwin'",
+        "tensorflow-text",
     ],
     extras_require={
         "extras": [