feat: Keras3 migration

.github/workflows/ci.yaml

@@ -11,94 +11,40 @@ env:
   UV_VERSION: 0.5.21
 
 jobs:
-  unittests_py_less_3_11:
-    name: Unit Tests Python=${{ matrix.python-version }} Pyspark=${{ matrix.pyspark-version }} Tensorflow=${{ matrix.tensorflow-version }}
+  unittests:
+    name: Unit Tests Python=${{ matrix.python-version }} Pyspark=${{ matrix.pyspark-version }} Keras=${{ matrix.keras-version }}
     runs-on: [ ubuntu-latest ]
     strategy:
       matrix:
-        # We match the last 2 Databricks LTS Runtime versions for pyspark
-        # and 3 Tensorflow versions within our package range that are not compatible with Python 3.11
-        python-version: ["3.9", "3.10"]
+        python-version: ["3.10", "3.11", "3.12"]
         pyspark-version: ["3.4.1", "3.5.0"]
-        tensorflow-version: ["2.9.1", "2.10.1", "2.11.1"]
+        keras-version: ["3.3.0", "3.7.0", "3.10.0", "3.12.0"]
     steps:
       - name: Setup Local Repo
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Setup Python
-        uses: actions/setup-python@v3
+        uses: actions/setup-python@v5
         with:
           python-version: ${{ matrix.python-version }}
       - name: Install uv
         run: |
           pip install --upgrade pip
           pip install "uv==$UV_VERSION"
-      - name: Fix Python Pyspark & Tensorflow Versions
+      - name: Install project and pin matrix versions
         run: |
           uv venv --python ${{ matrix.python-version }}
-          uv add pyspark==${{ matrix.pyspark-version }}
-          uv add tensorflow==${{ matrix.tensorflow-version }}
+          uv pip install -e ".[tensorflow]"
+          uv pip install pyspark==${{ matrix.pyspark-version }} keras==${{ matrix.keras-version }}
       - name: Run tests
         run: uv run -p ${{ matrix.python-version }} python -m pytest -n auto .
-  unittests_py_3_11:
-    name: Unit Tests Python=3.11 Pyspark=${{ matrix.pyspark-version }} Tensorflow=${{ matrix.tensorflow-version }}
-    runs-on: [ ubuntu-latest ]
-    strategy:
-      matrix:
-        # Only certain versions of pyspark and tensorflow are compatible with Python 3.11
-        pyspark-version: ["3.4.1", "3.5.0"]
-        tensorflow-version: ["2.12.1", "2.13.1", "2.14.1", "2.15.1", "2.16.2", "2.17.1", "2.18.0"]
-    steps:
-      - name: Setup Local Repo
-        uses: actions/checkout@v3
-      - name: Setup Python
-        uses: actions/setup-python@v3
-        with:
-          python-version: "3.11"
-      - name: Install uv
-        run: |
-          pip install --upgrade pip
-          pip install "uv==$UV_VERSION"
-      - name: Fix Python Pyspark & Tensorflow Versions
-        run: |
-          uv venv --python 3.11
-          uv add pyspark==${{ matrix.pyspark-version }}
-          uv add "tensorflow==${{ matrix.tensorflow-version }}; python_version >='3.9' and python_version <'3.12'"
-      - name: Run tests
-        run: uv run -p 3.11 python -m pytest -n auto .
-  unittests_py_3_12:
-    name: Unit Tests Python=3.12 Pyspark=${{ matrix.pyspark-version }} Tensorflow=${{ matrix.tensorflow-version }}
-    runs-on: [ ubuntu-latest ]
-    strategy:
-      matrix:
-        # Only certain versions of pyspark and tensorflow are compatible with Python 3.12
-        pyspark-version: [ "3.4.1", "3.5.0" ]
-        tensorflow-version: [ "2.16.2", "2.17.1", "2.18.0" ]
-    steps:
-      - name: Setup Local Repo
-        uses: actions/checkout@v3
-      - name: Setup Python
-        uses: actions/setup-python@v3
-        with:
-          python-version: "3.12"
-      - name: Install uv
-        run: |
-          pip install --upgrade pip
-          pip install "uv==$UV_VERSION"
-      - name: Fix Python Pyspark & Tensorflow Versions
-        run: |
-          uv venv --python 3.12
-          uv add pyspark==${{ matrix.pyspark-version }}
-          uv add "tensorflow==${{ matrix.tensorflow-version }}; python_version >='3.9' and python_version <='3.12'"
-      - name: Run tests
-        run: uv run -p 3.12 python -m pytest -n auto .
   formatting:
     name: Formatting Checks
     runs-on: [ ubuntu-latest ]
     steps:
       - name: Setup Local Repo
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Setup Python
-        uses: actions/setup-python@v3
+        uses: actions/setup-python@v5
         with:
           python-version: "3.10"
       - name: Install uv
@@ -112,9 +58,9 @@ jobs:
     runs-on: [ ubuntu-latest ]
     steps:
       - name: Setup Local Repo
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Setup Python
-        uses: actions/setup-python@v3
+        uses: actions/setup-python@v5
         with:
           python-version: "3.10"
       - name: Install uv
@@ -128,9 +74,9 @@ jobs:
     runs-on: [ ubuntu-latest ]
     steps:
       - name: Setup Local Repo
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Setup Python
-        uses: actions/setup-python@v3
+        uses: actions/setup-python@v5
         with:
           python-version: "3.10"
       - name: Install Pre-commit
@@ -145,9 +91,9 @@ jobs:
     runs-on: [ ubuntu-latest ]
     steps:
       - name: Setup Local Repo
-        uses: actions/checkout@v3
+        uses: actions/checkout@v4
       - name: Setup Python
-        uses: actions/setup-python@v3
+        uses: actions/setup-python@v5
         with:
           python-version: "3.10"
       - name: Install uv