update from keras-core to keras 3 (#970)

Author: haifeng-jin

.github/workflows/actions.yml

@@ -32,7 +32,8 @@ jobs:
           key: ${{ runner.os }}-pip-${{ hashFiles('setup.py') }}
       - name: Install dependencies
         run: |
-          pip install -e ".[tensorflow-cpu,tests]" --progress-bar off --upgrade
+          pip install -e ".[tests]" --progress-bar off --upgrade
+          pip install tensorflow-cpu==2.14.0
           pip install jax[cpu]
       - name: Test with pytest
         run: |
@@ -64,10 +65,11 @@ jobs:
           key: ${{ runner.os }}-pip-${{ hashFiles('setup.py') }}
       - name: Install dependencies
         run: |
-          pip install -e ".[tensorflow-cpu,tests]" --progress-bar off --upgrade
+          pip install -e ".[tests]" --progress-bar off --upgrade
           pip install torch>=2.0.1+cpu --progress-bar off
           pip install jax[cpu] --progress-bar off
-          pip install -e ".[tensorflow-cpu,tests]" --progress-bar off --upgrade
+          pip uninstall keras -y
+          pip install tf-nightly==2.16.0.dev20231103
       - name: Test with pytest
         env:
           KERAS_BACKEND: ${{ matrix.backend }}

keras_tuner/applications/efficientnet.py

@@ -16,7 +16,6 @@
 
 from keras_tuner.api_export import keras_tuner_export
 from keras_tuner.backend import keras
-from keras_tuner.backend import ops
 from keras_tuner.backend.keras import layers
 from keras_tuner.engine import hypermodel
 
@@ -121,12 +120,7 @@ def build(self, hp):
         )
         img_size = EFFICIENTNET_IMG_SIZE[version]
 
-        x = ops.image.resize(
-            x,
-            (img_size, img_size),
-            interpolation="bilinear",
-            data_format=keras.backend.image_data_format(),
-        )
+        x = layers.Resizing(img_size, img_size, interpolation="bilinear")(x)
         efficientnet_model = EFFICIENTNET_MODELS[version](
             include_top=False, input_tensor=x
         )

keras_tuner/backend/__init__.py

@@ -15,13 +15,12 @@
 Keras backend module.
 
 This module adds a temporarily Keras API surface that is fully under KerasTuner
-control. This allows us to switch between `keras_core` and `tf.keras`, as well
+control. This allows us to switch between Keras 3 and `tf.keras`, as well
 as add shims to support older version of `tf.keras`.
 
 - `config`: check which backend is being run.
-- `keras`: The full `keras` API (via `keras_core` or `tf.keras`).
-- `ops`: `keras_core.ops`, always tf backed if using `tf.keras`.
-- `random`: `keras_core.random`, always tf backed if using `tf.keras`.
+- `keras`: The full `keras` API (via `keras` 3 or `tf.keras`).
+- `ops`: `keras.ops`, always tf backed if using `tf.keras`.
 """
 
 from keras_tuner.backend import config

keras_tuner/backend/config.py

@@ -12,63 +12,22 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import json
-import os
+import keras
 
-import keras_core
 
-_MULTI_BACKEND = False
+def _multi_backend():
+    version_fn = getattr(keras, "version", None)
+    return version_fn and version_fn().startswith("3.")
 
-# Set Keras base dir path given KERAS_HOME env variable, if applicable.
-# Otherwise either ~/.keras or /tmp.
-if "KERAS_HOME" in os.environ:
-    _keras_dir = os.environ.get("KERAS_HOME")
-else:
-    _keras_base_dir = os.path.expanduser("~")
-    if not os.access(_keras_base_dir, os.W_OK):
-        _keras_base_dir = "/tmp"
-    _keras_dir = os.path.join(_keras_base_dir, ".keras")
 
-# Attempt to read KerasTuner config file.
-_config_path = os.path.expanduser(os.path.join(_keras_dir, "keras_tuner.json"))
-if os.path.exists(_config_path):
-    try:
-        with open(_config_path) as f:
-            _config = json.load(f)
-    except ValueError:
-        _config = {}
-    _MULTI_BACKEND = _config.get("multi_backend", _MULTI_BACKEND)
-
-# Save config file, if possible.
-if not os.path.exists(_keras_dir):
-    try:
-        os.makedirs(_keras_dir)
-    except OSError:
-        # Except permission denied and potential race conditions
-        # in multi-threaded environments.
-        pass
-
-if not os.path.exists(_config_path):
-    _config = {
-        "multi_backend": _MULTI_BACKEND,
-    }
-    try:
-        with open(_config_path, "w") as f:
-            f.write(json.dumps(_config, indent=4))
-    except IOError:
-        # Except permission denied.
-        pass
-
-# Use keras-core if KERAS_BACKEND is set in the environment.
-if "KERAS_BACKEND" in os.environ and os.environ["KERAS_BACKEND"]:
-    _MULTI_BACKEND = True
+_MULTI_BACKEND = _multi_backend()
 
 
 def multi_backend():
-    """Check if keras_core is enabled."""
+    """Check if multi-backend keras is enabled."""
     return _MULTI_BACKEND
 
 
 def backend():
     """Check the backend framework."""
-    return "tensorflow" if not multi_backend() else keras_core.config.backend()
+    return "tensorflow" if not multi_backend() else keras.config.backend()

keras_tuner/backend/keras.py

@@ -17,7 +17,7 @@
 from keras_tuner.backend.config import multi_backend
 
 if multi_backend():
-    from keras_core import *  # noqa: F403, F401
+    from keras import *  # noqa: F403, F401
 else:
     import tensorflow as tf
     from tensorflow.keras import *  # noqa: F403, F401

keras_tuner/backend/ops.py

@@ -12,41 +12,27 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-import keras_core
-
 from keras_tuner.backend.config import multi_backend
 
 if multi_backend():
-    from keras_core.src.ops import *  # noqa: F403, F401
+    from keras.src.ops import *  # noqa: F403, F401
 else:
-    from keras_core.src.backend.tensorflow import *  # noqa: F403, F401
-    from keras_core.src.backend.tensorflow.core import *  # noqa: F403, F401
-    from keras_core.src.backend.tensorflow.math import *  # noqa: F403, F401
-    from keras_core.src.backend.tensorflow.nn import *  # noqa: F403, F401
-    from keras_core.src.backend.tensorflow.numpy import *  # noqa: F403, F401
-
-
-if keras_core.config.backend() == "tensorflow" or not multi_backend():
+    import tensorflow as tf
+    from tensorflow import cast  # noqa: F403, F401
 
-    def take_along_axis(x, indices, axis=None):
-        import tensorflow as tf
+    def any_symbolic_tensors(args=None, kwargs=None):
+        args = args or ()
+        kwargs = kwargs or {}
+        for x in tf.nest.flatten((args, kwargs)):
+            if "KerasTensor" in x.__class__.__name__:
+                return True
+        return False
 
-        # TODO: move this workaround for dynamic shapes into keras-core.
-        if axis < 0:
-            axis = axis + indices.shape.rank
-        # If all shapes after axis are 1, squeeze them off and use tf.gather.
-        # tf.gather plays nicer with dynamic shapes in compiled functions.
-        leftover_axes = list(range(axis + 1, indices.shape.rank))
-        static_shape = indices.shape.as_list()
-        squeezable = True
-        for i in leftover_axes:
-            if static_shape[i] != 1:
-                squeezable = False
-        if squeezable:
-            if leftover_axes:
-                indices = tf.squeeze(indices, leftover_axes)
-            return tf.gather(x, indices, batch_dims=axis)
-        # Otherwise, fall back to the tfnp call.
-        return keras_core.src.backend.tensorflow.numpy.take_along_axis(
-            x, indices, axis=axis
+    def shape(x):
+        if any_symbolic_tensors((x,)):
+            return x.shape
+        dynamic = tf.shape(x)
+        static = x.shape.as_list()
+        return tuple(
+            dynamic[i] if s is None else s for i, s in enumerate(static)
         )

keras_tuner/backend/random.py

@@ -15,6 +15,6 @@
 from keras_tuner.backend.config import multi_backend
 
 if multi_backend():
-    from keras_core.random import *  # noqa: F403, F401
+    from keras.random import *  # noqa: F403, F401
 else:
-    from keras_core.src.backend.tensorflow.random import *  # noqa: F403, F401
+    from tensorflow.random import *  # noqa: F403, F401

setup.py

@@ -52,7 +52,7 @@ def get_version(rel_path):
     license="Apache License 2.0",
     version=VERSION,
     install_requires=[
-        "keras-core",
+        "keras",
         "packaging",
         "requests",
         "kt-legacy",