From 9c675a9a45e5e8244163fea82efc6066722608a1 Mon Sep 17 00:00:00 2001
From: Haifeng Jin <5476582+haifeng-jin@users.noreply.github.com>
Date: Mon, 27 Nov 2023 15:53:15 -0800
Subject: [PATCH] Bug fix for TF backend using random zoom with flatten
 (#18835)

* bug fix

* skip the numpy backend
---
 keras/backend/tensorflow/image.py              |  1 +
 keras/layers/preprocessing/random_zoom_test.py | 17 +++++++++++++++++
 2 files changed, 18 insertions(+)

diff --git a/keras/backend/tensorflow/image.py b/keras/backend/tensorflow/image.py
index 3fa90521e84..c6ccc06a850 100644
--- a/keras/backend/tensorflow/image.py
+++ b/keras/backend/tensorflow/image.py
@@ -119,6 +119,7 @@ def affine_transform(
         interpolation=interpolation.upper(),
         fill_mode=fill_mode.upper(),
     )
+    affined = tf.ensure_shape(affined, image.shape)
 
     if data_format == "channels_first":
         affined = tf.transpose(affined, (0, 3, 1, 2))
diff --git a/keras/layers/preprocessing/random_zoom_test.py b/keras/layers/preprocessing/random_zoom_test.py
index 21a3e87afdb..926f1951242 100644
--- a/keras/layers/preprocessing/random_zoom_test.py
+++ b/keras/layers/preprocessing/random_zoom_test.py
@@ -1,4 +1,5 @@
 import numpy as np
+import pytest
 from absl.testing import parameterized
 from tensorflow import data as tf_data
 
@@ -132,3 +133,19 @@ def test_dynamic_shape(self):
         )(inputs)
         model = models.Model(inputs, outputs)
         model.predict(np.random.random((1, 6, 6, 3)))
+
+    @pytest.mark.skipif(
+        backend.backend() == "numpy",
+        reason="The NumPy backend does not implement fit.",
+    )
+    def test_connect_with_flatten(self):
+        model = models.Sequential(
+            [
+                layers.RandomZoom((-0.5, 0.0), (-0.5, 0.0)),
+                layers.Flatten(),
+                layers.Dense(1, activation="relu"),
+            ],
+        )
+
+        model.compile(loss="mse")
+        model.fit(np.random.random((2, 2, 2, 1)), y=np.random.random((2,)))