Add support for multiple dynamic dimensions to Flatten layer on TF. (#21399)

With TensorFlow, when both the batch size and some other dimensions are dynamic, we need to use the dynamic batch size that comes from `ops.shape` instead of None / -1, otherwise there are two Nones / -1s passed to `ops.reshape`.

Fixes #21380

Author: hertschuh

keras/src/layers/reshaping/flatten.py

@@ -40,18 +40,22 @@ def __init__(self, data_format=None, **kwargs):
         self._channels_first = self.data_format == "channels_first"
 
     def call(self, inputs):
-        input_shape = inputs.shape
+        input_shape = ops.shape(inputs)
         rank = len(input_shape)
 
         if self._channels_first and rank > 1:
             # Switch to channels-last format.
             inputs = ops.transpose(inputs, axes=(0, *range(2, rank), 1))
 
-        output_shape = tuple(
-            dim if dim is not None else -1
-            for dim in self.compute_output_shape(input_shape)
-        )
-        return ops.reshape(inputs, output_shape)
+        non_batch_dims = input_shape[1:]
+        if len(non_batch_dims) == 0:
+            flattened_dim = 1
+        elif any(not isinstance(d, int) for d in non_batch_dims):
+            flattened_dim = -1
+        else:
+            flattened_dim = math.prod(non_batch_dims)
+
+        return ops.reshape(inputs, (input_shape[0], flattened_dim))
 
     def compute_output_shape(self, input_shape):
         non_batch_dims = input_shape[1:]

keras/src/layers/reshaping/flatten_test.py

@@ -2,8 +2,10 @@
 import pytest
 from absl.testing import parameterized
 
+from conftest import skip_if_backend
 from keras.src import backend
 from keras.src import layers
+from keras.src import models
 from keras.src import ops
 from keras.src import testing
 
@@ -112,12 +114,21 @@ def test_flatten_with_scalar_channels(self):
             expected_output=expected_output,
         )
 
-    def test_flatten_with_dynamic_batch_size(self):
+    def test_flatten_symbolic_with_dynamic_batch_size(self):
         input_layer = layers.Input(batch_shape=(None, 2, 3))
         flattened = layers.Flatten()(input_layer)
         self.assertEqual(flattened.shape, (None, 2 * 3))
 
-    def test_flatten_with_dynamic_dimension(self):
+    def test_flatten_symbolic_with_dynamic_dimension(self):
         input_layer = layers.Input(batch_shape=(5, 2, None))
         flattened = layers.Flatten()(input_layer)
         self.assertEqual(flattened.shape, (5, None))
+
+    @skip_if_backend("openvino", "Dynamic dimensions not supported by OpenVino")
+    def test_flatten_with_dynamic_batch_size_and_dynamic_dimenstions(self):
+        def generator():
+            yield (np.ones((3, 5, 7), dtype="float32"),)
+            yield (np.ones((2, 7, 5), dtype="float32"),)
+
+        model = models.Sequential([layers.Flatten()])
+        model.predict(generator())