Repository: keras-team/keras (python)

Issue:
apply_gradients AttributeError: 'ResourceVariable' object has no attribute 'overwrite_with_gradient'
When I have a mix of tf.Variable and KerasVariables I get the following error:
```
--> 632     if v.overwrite_with_gradient:
    633         if self.gradient_accumulation_steps:
    634             # Utilize a stateless manner for JAX compatibility
    635             steps = self.gradient_accumulation_steps

AttributeError: 'ResourceVariable' object has no attribute 'overwrite_with_gradient'
```

I suspect this is because my list of variables is [KerasVariables] + [tf.Variables]
and the following line only checks the first in the list as to whether overwrite_with_gradient can be used?
https://github.com/keras-team/keras/blob/660da946d33ccbb53575f15a17abdfd725cd2086/keras/src/optimizers/base_optimizer.py#L675


Pull request:
Fix `BaseOptimizer` to handle a mix of `tf.Variable` and `KerasVariable`
Fix #20517

Patch:
diff --git a/keras/src/callbacks/early_stopping.py b/keras/src/callbacks/early_stopping.py
index 5571cf606..bdf2112fc 100644
--- a/keras/src/callbacks/early_stopping.py
+++ b/keras/src/callbacks/early_stopping.py
@@ -136,14 +136,12 @@ class EarlyStopping(Callback):
             )
         if self.monitor_op == ops.less:
             self.min_delta *= -1
-        self.best = (
-            float("inf") if self.monitor_op == ops.less else -float("inf")
-        )
 
     def on_train_begin(self, logs=None):
         # Allow instances to be re-used
         self.wait = 0
         self.stopped_epoch = 0
+        self.best = None
         self.best_weights = None
         self.best_epoch = 0
 
@@ -210,4 +208,6 @@ class EarlyStopping(Callback):
         return monitor_value
 
     def _is_improvement(self, monitor_value, reference_value):
+        if reference_value is None:
+            return True
         return self.monitor_op(monitor_value - self.min_delta, reference_value)
diff --git a/keras/src/layers/preprocessing/image_preprocessing/random_zoom.py b/keras/src/layers/preprocessing/image_preprocessing/random_zoom.py
index 78473e4e2..9276d6923 100644
--- a/keras/src/layers/preprocessing/image_preprocessing/random_zoom.py
+++ b/keras/src/layers/preprocessing/image_preprocessing/random_zoom.py
@@ -3,7 +3,14 @@ from keras.src.api_export import keras_export
 from keras.src.layers.preprocessing.image_preprocessing.base_image_preprocessing_layer import (  # noqa: E501
     BaseImagePreprocessingLayer,
 )
+from keras.src.layers.preprocessing.image_preprocessing.bounding_boxes.converters import (  # noqa: E501
+    clip_to_image_size,
+)
+from keras.src.layers.preprocessing.image_preprocessing.bounding_boxes.converters import (  # noqa: E501
+    convert_format,
+)
 from keras.src.random.seed_generator import SeedGenerator
+from keras.src.utils import backend_utils
 
 
 @keras_export("keras.layers.RandomZoom")
@@ -175,13 +182,121 @@ class RandomZoom(BaseImagePreprocessingLayer):
     def transform_labels(self, labels, transformation, training=True):
         return labels
 
+    def get_transformed_x_y(self, x, y, transform):
+        a0, a1, a2, b0, b1, b2, c0, c1 = self.backend.numpy.split(
+            transform, 8, axis=-1
+        )
+
+        k = c0 * x + c1 * y + 1
+        x_transformed = (a0 * x + a1 * y + a2) / k
+        y_transformed = (b0 * x + b1 * y + b2) / k
+        return x_transformed, y_transformed
+
+    def get_clipped_bbox(self, bounding_boxes, h_end, h_start, w_end, w_start):
+        bboxes = bounding_boxes["boxes"]
+        x1, y1, x2, y2 = self.backend.numpy.split(bboxes, 4, axis=-1)
+
+        if len(bboxes.shape) == 3:
+            h_end = self.backend.numpy.expand_dims(h_end, -1)
+            h_start = self.backend.numpy.expand_dims(h_start, -1)
+            w_end = self.backend.numpy.expand_dims(w_end, -1)
+            w_start = self.backend.numpy.expand_dims(w_start, -1)
+
+        x1 = self.backend.numpy.clip(x1, w_start, w_end) - w_start
+        y1 = self.backend.numpy.clip(y1, h_start, h_end) - h_start
+        x2 = self.backend.numpy.clip(x2, w_start, w_end) - w_start
+        y2 = self.backend.numpy.clip(y2, h_start, h_end) - h_start
+        bounding_boxes["boxes"] = self.backend.numpy.concatenate(
+            [x1, y1, x2, y2], axis=-1
+        )
+        return bounding_boxes
+
     def transform_bounding_boxes(
         self,
         bounding_boxes,
         transformation,
         training=True,
     ):
-        raise NotImplementedError
+        if backend_utils.in_tf_graph():
+            self.backend.set_backend("tensorflow")
+
+        width_zoom = transformation["width_zoom"]
+        height_zoom = transformation["height_zoom"]
+        inputs_shape = transformation["input_shape"]
+
+        if self.data_format == "channels_first":
+            height = inputs_shape[-2]
+            width = inputs_shape[-1]
+        else:
+            height = inputs_shape[-3]
+            width = inputs_shape[-2]
+
+        bounding_boxes = convert_format(
+            bounding_boxes,
+            source=self.bounding_box_format,
+            target="xyxy",
+            height=height,
+            width=width,
+        )
+
+        zooms = self.backend.cast(
+            self.backend.numpy.concatenate([width_zoom, height_zoom], axis=1),
+            dtype="float32",
+        )
+        transform = self._get_zoom_matrix(zooms, height, width)
+
+        w_start, h_start = self.get_transformed_x_y(
+            0,
+            0,
+            transform,
+        )
+
+        w_end, h_end = self.get_transformed_x_y(
+            width,
+            height,
+            transform,
+        )
+
+        bounding_boxes = self.get_clipped_bbox(
+            bounding_boxes, h_end, h_start, w_end, w_start
+        )
+
+        height_transformed = h_end - h_start
+        width_transformed = w_end - w_start
+
+        height_transformed = self.backend.numpy.expand_dims(
+            height_transformed, -1
+        )
+        width_transformed = self.backend.numpy.expand_dims(
+            width_transformed, -1
+        )
+
+        bounding_boxes = convert_format(
+            bounding_boxes,
+            source="xyxy",
+            target="rel_xyxy",
+            height=height_transformed,
+            width=width_transformed,
+        )
+
+        bounding_boxes = clip_to_image_size(
+            bounding_boxes=bounding_boxes,
+            height=height_transformed,
+            width=width_transformed,
+            format="rel_xyxy",
+        )
+
+        bounding_boxes = convert_format(
+            bounding_boxes,
+            source="rel_xyxy",
+            target=self.bounding_box_format,
+            height=height,
+            width=width,
+        )
+
+        self.backend.reset()
+
+        return bounding_boxes
 
     def transform_segmentation_masks(
         self, segmentation_masks, transformation, training=True
@@ -229,6 +344,7 @@ class RandomZoom(BaseImagePreprocessingLayer):
         return {
             "height_zoom": height_zoom,
             "width_zoom": width_zoom,
+            "input_shape": images_shape,
         }
 
     def _zoom_inputs(self, inputs, transformation):
diff --git a/keras/src/optimizers/base_optimizer.py b/keras/src/optimizers/base_optimizer.py
index b966c15bc..8717192fa 100644
--- a/keras/src/optimizers/base_optimizer.py
+++ b/keras/src/optimizers/base_optimizer.py
@@ -672,7 +672,7 @@ class BaseOptimizer(KerasSaveable):
         """
         # Shortcut for `tf.Variable` because it doesn't have a
         # `overwrite_with_gradient` attr
-        if not hasattr(vars[0], "overwrite_with_gradient"):
+        if any(not hasattr(v, "overwrite_with_gradient") for v in vars):
             return grads, vars
 
         # Shallow copies
@@ -723,7 +723,11 @@ class BaseOptimizer(KerasSaveable):
             if filtered_grads[i] is None:
                 filtered_grads.pop(i)
                 v = filtered_vars.pop(i)
-                missing_grad_vars.append(v.path)
+                try:
+                    missing_grad_vars.append(v.path)
+                except AttributeError:
+                    # `tf.Variable` doesn't have `path` attr.
+                    missing_grad_vars.append(v.name)
 
         if not filtered_grads:
             raise ValueError("No gradients provided for any variable.")


Test patch:
diff --git a/keras/src/callbacks/early_stopping_test.py b/keras/src/callbacks/early_stopping_test.py
index e120fe8a2..d4b127675 100644
--- a/keras/src/callbacks/early_stopping_test.py
+++ b/keras/src/callbacks/early_stopping_test.py
@@ -114,16 +114,17 @@ class EarlyStoppingTest(testing.TestCase):
             loss="mae",
             metrics=["mse"],
         )
-        weights = model.get_weights()
+        stopper = callbacks.EarlyStopping(monitor="mse", patience=patience)
 
-        # This should allow training to go for at least `patience` epochs
-        model.set_weights(weights)
+        history1 = model.fit(
+            data, labels, callbacks=[stopper], verbose=0, epochs=20
+        )
+        self.assertGreaterEqual(len(history1.epoch), patience)
 
-        stopper = callbacks.EarlyStopping(monitor="mse", patience=patience)
-        hist = model.fit(
+        history2 = model.fit(
             data, labels, callbacks=[stopper], verbose=0, epochs=20
         )
-        assert len(hist.epoch) >= patience
+        self.assertGreaterEqual(len(history2.epoch), patience)
 
     @pytest.mark.requires_trainable_backend
     def test_early_stopping_with_baseline(self):
diff --git a/keras/src/layers/preprocessing/image_preprocessing/random_zoom_test.py b/keras/src/layers/preprocessing/image_preprocessing/random_zoom_test.py
index 52e084924..88be12f95 100644
--- a/keras/src/layers/preprocessing/image_preprocessing/random_zoom_test.py
+++ b/keras/src/layers/preprocessing/image_preprocessing/random_zoom_test.py
@@ -7,6 +7,7 @@ from keras.src import backend
 from keras.src import layers
 from keras.src import models
 from keras.src import testing
+from keras.src.utils import backend_utils
 
 
 class RandomZoomTest(testing.TestCase):
@@ -148,3 +149,121 @@ class RandomZoomTest(testing.TestCase):
 
         model.compile(loss="mse")
         model.fit(np.random.random((2, 2, 2, 1)), y=np.random.random((2,)))
+
+    @parameterized.named_parameters(
+        (
+            "with_zoom_in",
+            [[[0.1]], [[0.1]]],
+            [[[0.0, 0.0, 8.0, 0.0], [8.0, 0.0, 8.0, 10.0]]],
+        ),
+        (
+            "with_zoom_out",
+            [[[1.9]], [[1.9]]],
+            [
+                [
+                    [2.710526, 2.657895, 3.763158, 3.710526],
+                    [4.815789, 4.236842, 5.868421, 5.289474],
+                ]
+            ],
+        ),
+    )
+    def test_random_flip_bounding_boxes(self, zoom, expected_boxes):
+        data_format = backend.config.image_data_format()
+        if data_format == "channels_last":
+            image_shape = (10, 8, 3)
+        else:
+            image_shape = (3, 10, 8)
+        input_image = np.random.random(image_shape)
+        bounding_boxes = {
+            "boxes": np.array(
+                [
+                    [2, 1, 4, 3],
+                    [6, 4, 8, 6],
+                ]
+            ),
+            "labels": np.array([[1, 2]]),
+        }
+        input_data = {"images": input_image, "bounding_boxes": bounding_boxes}
+        random_zoom_layer = layers.RandomZoom(
+            height_factor=(0.5, 0.5),
+            data_format=data_format,
+            seed=42,
+            bounding_box_format="xyxy",
+        )
+
+        transformation = {
+            "height_zoom": backend_utils.convert_tf_tensor(np.array(zoom[0])),
+            "width_zoom": backend_utils.convert_tf_tensor(np.array(zoom[1])),
+            "input_shape": image_shape,
+        }
+        output = random_zoom_layer.transform_bounding_boxes(
+            input_data["bounding_boxes"],
+            transformation=transformation,
+            training=True,
+        )
+
+        self.assertAllClose(output["boxes"], expected_boxes)
+
+    @parameterized.named_parameters(
+        (
+            "with_zoom_in",
+            [[[0.1]], [[0.1]]],
+            [[[0.0, 0.0, 8.0, 0.0], [8.0, 0.0, 8.0, 10.0]]],
+        ),
+        (
+            "with_zoom_out",
+            [[[1.9]], [[1.9]]],
+            [
+                [
+                    [2.710526, 2.657895, 3.763158, 3.710526],
+                    [4.815789, 4.236842, 5.868421, 5.289474],
+                ]
+            ],
+        ),
+    )
+    def test_random_flip_tf_data_bounding_boxes(self, zoom, expected_boxes):
+        data_format = backend.config.image_data_format()
+        if backend.config.image_data_format() == "channels_last":
+            image_shape = (1, 10, 8, 3)
+        else:
+            image_shape = (1, 3, 10, 8)
+        input_image = np.random.random(image_shape)
+        bounding_boxes = {
+            "boxes": np.array(
+                [
+                    [
+                        [2, 1, 4, 3],
+                        [6, 4, 8, 6],
+                    ]
+                ]
+            ),
+            "labels": np.array([[1, 2]]),
+        }
+
+        input_data = {"images": input_image, "bounding_boxes": bounding_boxes}
+
+        ds = tf_data.Dataset.from_tensor_slices(input_data)
+        random_zoom_layer = layers.RandomZoom(
+            height_factor=0.5,
+            data_format=data_format,
+            seed=42,
+            bounding_box_format="xyxy",
+        )
+
+        transformation = {
+            "height_zoom": backend_utils.convert_tf_tensor(np.array(zoom[0])),
+            "width_zoom": backend_utils.convert_tf_tensor(np.array(zoom[1])),
+            "input_shape": image_shape,
+        }
+
+        ds = ds.map(
+            lambda x: random_zoom_layer.transform_bounding_boxes(
+                x["bounding_boxes"],
+                transformation=transformation,
+                training=True,
+            )
+        )
+
+        output = next(iter(ds))
+        expected_boxes = np.array(expected_boxes)
+        self.assertAllClose(output["boxes"], expected_boxes)
diff --git a/keras/src/optimizers/optimizer_test.py b/keras/src/optimizers/optimizer_test.py
index 1ac35e63c..7d661df9a 100644
--- a/keras/src/optimizers/optimizer_test.py
+++ b/keras/src/optimizers/optimizer_test.py
@@ -401,3 +401,25 @@ class OptimizerTest(testing.TestCase):
         reloaded = pickle.loads(pickle.dumps(optimizer))
 
         self.assertEqual(optimizer.get_config(), reloaded.get_config())
+
+    @pytest.mark.skipif(
+        backend.backend() != "tensorflow",
+        reason="The tf.Variable test can only run with TensorFlow backend.",
+    )
+    def test_mixed_with_tf_variables(self):
+        import tensorflow as tf
+
+        v = backend.Variable([[1.0, 2.0], [3.0, 4.0]])
+        grads = backend.convert_to_tensor([[1.0, 1.0], [1.0, 1.0]])
+        tf_v = tf.Variable([[1.0, 2.0], [3.0, 4.0]])
+        tf_grads = backend.convert_to_tensor([[1.0, 1.0], [1.0, 1.0]])
+        optimizer = optimizers.Adam(learning_rate=1.0)
+        optimizer.apply_gradients([(grads, v), (tf_grads, tf_v)])
+        self.assertAllClose(optimizer.iterations, 1)
+
+        # Test with no grads
+        with self.assertWarnsRegex(
+            UserWarning, "Gradients do not exist for variables"
+        ):
+            optimizer.apply_gradients([(grads, v), (None, tf_v)])
+            self.assertAllClose(optimizer.iterations, 2)