Ditch Keras 2 support (#1658)

* Ditch Keras 2

* Try workaround for tf take_along_axis error

* Remove dependencies

.github/workflows/actions.yml

@@ -11,40 +11,8 @@ permissions:
   contents: read
 
 jobs:
-  keras_2:
-    name: Test the code with Keras 2
-    runs-on: ubuntu-latest
-    steps:
-      - uses: actions/checkout@v4
-      - name: Set up Python 3.9
-        uses: actions/setup-python@v5
-        with:
-          python-version: 3.9
-      - name: Get pip cache dir
-        id: pip-cache
-        run: |
-          python -m pip install --upgrade pip setuptools
-          echo "::set-output name=dir::$(pip cache dir)"
-      - name: pip cache
-        uses: actions/cache@v4
-        with:
-          path: ${{ steps.pip-cache.outputs.dir }}
-          key: ${{ runner.os }}-pip-${{ hashFiles('setup.py') }}
-          restore-keys: |
-            ${{ runner.os }}-pip-
-      - name: Install dependencies
-        run: |
-          pip install -r requirements-common.txt --progress-bar off
-          pip install tensorflow-text==2.14 tensorflow==2.14 keras-core
-          pip install --no-deps -e "." --progress-bar off
-      - name: Test with pytest
-        run: |
-          pytest --cov=keras_nlp --cov-report xml:coverage.xml keras_nlp/
-      - name: Run integration tests
-        run: |
-          python pip_build.py --install && cd integration_tests && pytest .
-  keras_3:
-    name: Test the code with Keras 3
+  run_tests:
+    name: Test the code
     strategy:
       fail-fast: false
       matrix:

conftest.py

@@ -24,8 +24,7 @@
         "The TensorFlow package is required for data preprocessing with any backend."
     )
 
-from keras_nlp.src.backend import config as backend_config
-from keras_nlp.src.backend import keras
+import keras
 
 
 def pytest_addoption(parser):
@@ -59,7 +58,7 @@ def pytest_configure(config):
     # Verify that device has GPU and detected by backend
     if config.getoption("--check_gpu"):
         found_gpu = False
-        backend = backend_config.backend()
+        backend = keras.config.backend()
         if backend == "jax":
             import jax
 
@@ -88,10 +87,6 @@ def pytest_configure(config):
         "markers",
         "tf_only: mark test as a tf only test",
     )
-    config.addinivalue_line(
-        "markers",
-        "keras_3_only: mark test as a keras 3 only test",
-    )
     config.addinivalue_line(
         "markers",
         "kaggle_key_required: mark test needing a kaggle key",
@@ -113,13 +108,9 @@ def pytest_collection_modifyitems(config, items):
         reason="need --run_extra_large option to run",
     )
     tf_only = pytest.mark.skipif(
-        not backend_config.backend() == "tensorflow",
+        not keras.config.backend() == "tensorflow",
         reason="tests only run on tf backend",
     )
-    keras_3_only = pytest.mark.skipif(
-        not backend_config.keras_3(),
-        reason="tests only run on with multi-backend keras",
-    )
     found_kaggle_key = all(
         [
             os.environ.get("KAGGLE_USERNAME", None),
@@ -137,13 +128,9 @@ def pytest_collection_modifyitems(config, items):
             item.add_marker(skip_extra_large)
         if "tf_only" in item.keywords:
             item.add_marker(tf_only)
-        if "keras_3_only" in item.keywords:
-            item.add_marker(keras_3_only)
         if "kaggle_key_required" in item.keywords:
             item.add_marker(kaggle_key_required)
 
 
 # Disable traceback filtering for quicker debugging of tests failures.
-tf.debugging.disable_traceback_filtering()
-if backend_config.keras_3():
-    keras.config.disable_traceback_filtering()
+keras.config.disable_traceback_filtering()

keras_nlp/src/api_export.py

@@ -14,7 +14,7 @@
 
 import types
 
-from keras_nlp.src.backend import keras
+import keras
 
 try:
     import namex

keras_nlp/src/layers/modeling/alibi_bias.py

@@ -13,9 +13,10 @@
 # limitations under the License.
 import math
 
+import keras
+from keras import ops
+
 from keras_nlp.src.api_export import keras_nlp_export
-from keras_nlp.src.backend import keras
-from keras_nlp.src.backend import ops
 
 
 @keras_nlp_export("keras_nlp.layers.AlibiBias")

keras_nlp/src/layers/modeling/alibi_bias_test.py

@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from keras_nlp.src.backend import keras
-from keras_nlp.src.backend import ops
-from keras_nlp.src.backend import random
+import keras
+from keras import ops
+from keras import random
+
 from keras_nlp.src.layers.modeling.alibi_bias import AlibiBias
 from keras_nlp.src.tests.test_case import TestCase
 

keras_nlp/src/layers/modeling/cached_multi_head_attention.py

@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import keras
+from keras import ops
+
 from keras_nlp.src.api_export import keras_nlp_export
-from keras_nlp.src.backend import keras
-from keras_nlp.src.backend import ops
 
 
 @keras_nlp_export("keras_nlp.layers.CachedMultiHeadAttention")

keras_nlp/src/layers/modeling/cached_multi_head_attention_test.py

@@ -12,9 +12,9 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-from keras_nlp.src.backend import config
-from keras_nlp.src.backend import ops
-from keras_nlp.src.backend import random
+from keras import ops
+from keras import random
+
 from keras_nlp.src.layers.modeling.cached_multi_head_attention import (
     CachedMultiHeadAttention,
 )
@@ -37,9 +37,6 @@ def test_layer_behaviors(self):
             expected_output_shape=(2, 4, 6),
             expected_num_trainable_weights=8,
             expected_num_non_trainable_variables=1,
-            # Keras 2 does not handle mixed precision correctly when not set
-            # globally.
-            run_precision_checks=config.keras_3(),
         )
 
     def test_cache_call_is_correct(self):

keras_nlp/src/layers/modeling/f_net_encoder.py

@@ -12,9 +12,10 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
+import keras
+from keras import ops
+
 from keras_nlp.src.api_export import keras_nlp_export
-from keras_nlp.src.backend import keras
-from keras_nlp.src.backend import ops
 from keras_nlp.src.utils.keras_utils import clone_initializer