Tf generator changes (#2162)

* Keras3 tests (#2152)

* Update requirements for Keras 3

* Update Requirements

* Update Requirements

* Add randint for keras3 compatibility

* remove force generator

* fix few tests

* fix tests

* fix more tests

* Revert "Replace `RandomGenerator` with `SeedGenerator` (#2150)" (#2161)

This reverts commit 365a675.

* Update actions.yml

---------

Co-authored-by: Ramesh Sampath <1437573+sampathweb@users.noreply.github.com>

.github/workflows/actions.yml

@@ -11,13 +11,13 @@ permissions:
   contents: read
 
 jobs:
-  keras2:
-    name: Test the code with tf.keras
+  keras_2:
+    name: Test the code with Keras 2
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
     - name: Set up Python 3.9
-      uses: actions/setup-python@v1
+      uses: actions/setup-python@v4
       with:
         python-version: 3.9
     - name: Get pip cache dir
@@ -34,7 +34,7 @@ jobs:
           ${{ runner.os }}-pip-
     - name: Install dependencies
       run: |
-        pip install tensorflow>=2.13.0
+        pip install tensorflow~=2.14
         pip install torch>=2.0.1+cpu
         pip install "jax[cpu]"
         pip install keras-core
@@ -44,17 +44,17 @@ jobs:
         TEST_CUSTOM_OPS: false
       run: |
         pytest keras_cv/ --ignore keras_cv/models/legacy/ --durations 0
-  multibackend:
-    name: Test the code with Keras Core
+  keras_3:
+    name: Test the code with Keras 3
     strategy:
       fail-fast: false
       matrix:
         backend: [tensorflow, jax, torch]
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
     - name: Set up Python 3.9
-      uses: actions/setup-python@v1
+      uses: actions/setup-python@v4
       with:
         python-version: 3.9
     - name: Get pip cache dir
@@ -71,11 +71,7 @@ jobs:
           ${{ runner.os }}-pip-
     - name: Install dependencies
       run: |
-        pip install tensorflow>=2.13.0
-        pip install "jax[cpu]"
-        pip install torch>=2.0.1+cpu
-        pip install torchvision>=0.15.1
-        pip install keras-core
+        pip install -r requirements.txt
         pip install -e ".[tests]" --progress-bar off --upgrade
     - name: Test with pytest
       env:
@@ -99,9 +95,9 @@ jobs:
     name: Check the code format
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
+    - uses: actions/checkout@v3
     - name: Set up Python 3.9
-      uses: actions/setup-python@v1
+      uses: actions/setup-python@v4
       with:
         python-version: 3.9
     - name: Get pip cache dir
@@ -118,7 +114,7 @@ jobs:
           ${{ runner.os }}-pip-
     - name: Install dependencies
       run: |
-        pip install tensorflow>=2.13.0
+        pip install tensorflow~=2.14
         pip install -e ".[tests]" --progress-bar off --upgrade
     - name: Lint
       run: bash shell/lint.sh

.github/workflows/devcontainer.yml

@@ -20,7 +20,7 @@ jobs:
     steps:
 
       - name: Checkout (GitHub)
-        uses: actions/checkout@v2
+        uses: actions/checkout@v4
 
       - name: Build and run dev container task
         uses: devcontainers/ci@v0.2

.github/workflows/nightly.yml

@@ -40,6 +40,7 @@ jobs:
           pip install -r requirements.txt --progress-bar off
       - name: Build wheel file
         run: |
+          export BUILD_WITH_CUSTOM_OPS=false
           python pip_build.py --nightly
       - name: Publish to PyPI
         uses: pypa/gh-action-pypi-publish@release/v1

.github/workflows/release.yml

@@ -8,105 +8,38 @@ permissions:
   contents: read
 
 jobs:
-  deploy-with-custom-ops:
-    # This job is currently skipped until we cut a release with custom ops.
-    if: false
-    name: Build and deploy release wheels with custom ops
-    runs-on: ${{ matrix.os }}
-    strategy:
-      matrix:
-        # To switch on windows-2022/latest, please verify the bazel version:
-        # https://github.com/bazelbuild/bazel/issues/14232#issuecomment-1011247429
-        os: ['macos-12', 'windows-2019', 'ubuntu-18.04']
-        py-version: ['3.9', '3.10', '3.11']
-        tf-version: ['2.13.0']
-        use-macos-arm: [false]
-        include:
-          - os: 'macos-12'
-            tf-version: '2.13.0'
-            py-version: '3.9'
-            use-macos-arm: true
-          - os: 'macos-12'
-            tf-version: '2.13.0'
-            py-version: '3.10'
-            use-macos-arm: true
-      fail-fast: false
-    steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python
-      uses: actions/setup-python@v1
-      with:
-        python-version: ${{ matrix.py-version }}
-    - name: Get pip cache dir
-      id: pip-cache
-      run: |
-        python -m pip install --upgrade pip setuptools wheel auditwheel twine
-        echo "::set-output name=dir::$(pip cache dir)"
-    - name: pip cache
-      uses: actions/cache@v2
-      with:
-        path: ${{ steps.pip-cache.outputs.dir }}
-        key: ${{ runner.os }}-pip-${{ hashFiles('setup.py') }}
-        restore-keys: |
-          ${{ runner.os }}-pip-
-    - name: Install dependencies
-      run: |
-        python -m pip install tensorflow-cpu==${{ matrix.tf-version }}
-        python -m pip install -e ".[tests]" --progress-bar off --upgrade
-    - name: Configure Build Environment
-      run: |
-        python build_deps/configure.py
-    - name: Reinstall TensorFlow (MacOS ARM)
-      if: ${{ matrix.os == 'macos-12' && matrix.use-macos-arm}}
-      run: |
-        python -m pip uninstall -y tensorflow-cpu
-        python -m pip install --platform=macosx_12_0_arm64 --no-deps --target=$(python -c 'import site; print(site.getsitepackages()[0])') --upgrade tensorflow-macos==${{ matrix.tf-version }}
-    - name: Bazel Build
-      if: ${{ ! matrix.use-macos-arm }}
-      run: |
-        export BUILD_WITH_CUSTOM_OPS=true
-        bazel build build_pip_pkg
-    - name: Bazel Build (MacOS ARM)
-      if: ${{ matrix.use-macos-arm}}
-      run: |
-        bazel build --cpu=darwin_arm64 --copt -mmacosx-version-min=12.0 --linkopt -mmacosx-version-min=12.0 build_pip_pkg
-    - name: Build wheels
-      run: |
-        export BUILD_WITH_CUSTOM_OPS=true
-        bazel-bin/build_pip_pkg wheels
-    - name: Repair wheels (manylinux)
-      if: ${{ matrix.os == 'ubuntu-18.04' }}
-      run: |
-        python -m pip install --upgrade patchelf==0.14
-        bash build_deps/tf_auditwheel_patch.sh
-        python -m auditwheel repair --plat manylinux2014_x86_64 wheels/*.whl
-        rm wheels/*.whl
-        mv wheelhouse/* wheels/
-    - name: Upload wheels
-      env:
-        TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
-        TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
-      run: |
-        twine upload wheels/*
-  deploy-without-custom-ops:
-    name: Build and deploy release wheels without custom ops
+  run-test-for-release:
+    uses: ./.github/workflows/actions.yml
+  release:
+    name: Build and publish to PyPI
     runs-on: ubuntu-latest
     steps:
-    - uses: actions/checkout@v2
-    - name: Set up Python
-      uses: actions/setup-python@v1
-      with:
-        python-version: 3.9
-    - name: Build wheels
-      run: |
-        pip install tensorflow==2.13.0
-        python -m pip install --upgrade setuptools wheel twine
-        python -m pip install --upgrade -r requirements.txt
-        export BUILD_WITH_CUSTOM_OPS=false
-        python pip_build.py
-    - name: Upload wheels
-      env:
-        TWINE_USERNAME: ${{ secrets.PYPI_USERNAME }}
-        TWINE_PASSWORD: ${{ secrets.PYPI_PASSWORD }}
-      run: |
-        twine upload dist/*.whl
+      - uses: actions/checkout@v3
+      - name: Set up Python
+        uses: actions/setup-python@v4
+        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@v2
+        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.txt --progress-bar off
+      - name: Build a binary wheel and a source tarball
+        run: |
+          export BUILD_WITH_CUSTOM_OPS=false
+          python pip_build.py
+      - name: Publish distribution to PyPI
+        if: startsWith(github.ref, 'refs/tags')
+        uses: pypa/gh-action-pypi-publish@master
+        with:
+          password: ${{ secrets.PYPI_API_TOKEN }}

benchmarks/vectorized_jittered_resize.py

@@ -20,15 +20,14 @@
 from tensorflow import keras
 
 from keras_cv import bounding_box
-from keras_cv.backend import random
 from keras_cv.layers import JitteredResize
 from keras_cv.layers.preprocessing.base_image_augmentation_layer import (
     BaseImageAugmentationLayer,
 )
 from keras_cv.layers.preprocessing.vectorized_base_image_augmentation_layer import (  # noqa: E501
     BOUNDING_BOXES,
 )
-from keras_cv.layers.preprocessing.vectorized_base_image_augmentation_layer import (  # noqa: E501
+from keras_cv.layers.preprocessing.vectorized_base_image_augmentation_layer import (
     IMAGES,
 )
 from keras_cv.utils import preprocessing as preprocessing_utils
@@ -259,7 +258,7 @@ def test_consistency_with_old_impl(self):
 
         # makes offsets fixed to (0.5, 0.5)
         with unittest.mock.patch.object(
-            random,
+            layer._random_generator,
             "uniform",
             return_value=tf.convert_to_tensor([[0.5, 0.5]]),
         ):

benchmarks/vectorized_mosaic.py

@@ -20,15 +20,14 @@
 from tensorflow import keras
 
 from keras_cv import bounding_box
-from keras_cv.backend import random
 from keras_cv.layers import Mosaic
 from keras_cv.layers.preprocessing.base_image_augmentation_layer import (
     BaseImageAugmentationLayer,
 )
 from keras_cv.layers.preprocessing.vectorized_base_image_augmentation_layer import (  # noqa: E501
     IMAGES,
 )
-from keras_cv.layers.preprocessing.vectorized_base_image_augmentation_layer import (  # noqa: E501
+from keras_cv.layers.preprocessing.vectorized_base_image_augmentation_layer import (
     LABELS,
 )
 from keras_cv.utils import preprocessing as preprocessing_utils
@@ -102,7 +101,9 @@ def _batch_augment(self, inputs):
             minval=0,
             maxval=batch_size,
             dtype=tf.int32,
-            seed=random.make_seed(seed=self._seed_generator),
+            seed=self._random_generator.uniform(
+                shape=[], minval=0, maxval=9999, dtype="int32"
+            ),
         )
         # concatenate the batches with permutation order to get all 4 images of
         # the mosaic

benchmarks/vectorized_random_brightness.py

@@ -61,7 +61,7 @@ class OldRandomBrightness(BaseImageAugmentationLayer):
     """
 
     def __init__(self, factor, value_range=(0, 255), seed=None, **kwargs):
-        super().__init__(seed=seed, force_generator=True, **kwargs)
+        super().__init__(seed=seed, **kwargs)
         if isinstance(factor, float) or isinstance(factor, int):
             factor = (-factor, factor)
         self.factor = preprocessing_utils.parse_factor(

benchmarks/vectorized_random_contrast.py

@@ -60,7 +60,7 @@ class OldRandomContrast(BaseImageAugmentationLayer):
     """
 
     def __init__(self, factor, seed=None, **kwargs):
-        super().__init__(seed=seed, force_generator=True, **kwargs)
+        super().__init__(seed=seed, **kwargs)
         if isinstance(factor, (tuple, list)):
             min = 1 - factor[0]
             max = 1 + factor[1]

benchmarks/vectorized_random_crop.py

@@ -21,7 +21,6 @@
 from tensorflow import keras
 
 from keras_cv import bounding_box
-from keras_cv.backend import random
 from keras_cv.layers import RandomCrop
 from keras_cv.layers.preprocessing.base_image_augmentation_layer import (
     BaseImageAugmentationLayer,
@@ -60,7 +59,9 @@ def __init__(
         self, height, width, seed=None, bounding_box_format=None, **kwargs
     ):
         super().__init__(
-            **kwargs, autocast=False, seed=seed, force_generator=True
+            **kwargs,
+            autocast=False,
+            seed=seed,
         )
         self.height = height
         self.width = width
@@ -73,9 +74,7 @@ def get_random_transformation(self, image=None, **kwargs):
         h_diff = image_shape[H_AXIS] - self.height
         w_diff = image_shape[W_AXIS] - self.width
         dtype = image_shape.dtype
-        rands = random.uniform(
-            [2], 0, dtype.max, dtype, seed=self._seed_generator
-        )
+        rands = self._random_generator.uniform([2], 0, dtype.max, dtype)
         h_start = rands[0] % (h_diff + 1)
         w_start = rands[1] % (w_diff + 1)
         return {"top": h_start, "left": w_start}

benchmarks/vectorized_random_flip.py

@@ -20,15 +20,14 @@
 from tensorflow import keras
 
 from keras_cv import bounding_box
-from keras_cv.backend import random
 from keras_cv.layers import RandomFlip
 from keras_cv.layers.preprocessing.base_image_augmentation_layer import (
     BaseImageAugmentationLayer,
 )
 from keras_cv.layers.preprocessing.vectorized_base_image_augmentation_layer import (  # noqa: E501
     BOUNDING_BOXES,
 )
-from keras_cv.layers.preprocessing.vectorized_base_image_augmentation_layer import (  # noqa: E501
+from keras_cv.layers.preprocessing.vectorized_base_image_augmentation_layer import (
     IMAGES,
 )
 
@@ -72,7 +71,7 @@ class OldRandomFlip(BaseImageAugmentationLayer):
     def __init__(
         self, mode=HORIZONTAL, seed=None, bounding_box_format=None, **kwargs
     ):
-        super().__init__(seed=seed, force_generator=True, **kwargs)
+        super().__init__(seed=seed, **kwargs)
         self.mode = mode
         self.seed = seed
         if mode == HORIZONTAL:
@@ -102,13 +101,9 @@ def get_random_transformation(self, **kwargs):
         flip_horizontal = False
         flip_vertical = False
         if self.horizontal:
-            flip_horizontal = (
-                random.uniform(shape=[], seed=self._seed_generator) > 0.5
-            )
+            flip_horizontal = self._random_generator.uniform(shape=[]) > 0.5
         if self.vertical:
-            flip_vertical = (
-                random.uniform(shape=[], seed=self._seed_generator) > 0.5
-            )
+            flip_vertical = self._random_generator.uniform(shape=[]) > 0.5
         return {
             "flip_horizontal": tf.cast(flip_horizontal, dtype=tf.bool),
             "flip_vertical": tf.cast(flip_vertical, dtype=tf.bool),
@@ -237,13 +232,13 @@ def test_consistency_with_old_impl(self):
         )
 
         with unittest.mock.patch.object(
-            random,
+            layer._random_generator,
             "uniform",
             return_value=tf.convert_to_tensor([[0.6]]),
         ):
             output = layer(image)
         with unittest.mock.patch.object(
-            random,
+            old_layer._random_generator,
             "uniform",
             return_value=tf.convert_to_tensor(0.6),
         ):