Proposed fix for issue keras-team#21519: Reshape layer does not handle -1 shape… (keras-team#21568)

* Proposed fix for issue keras-team#21519: Reshape layer does not handle -1 shape infor dynamically.

* Removed unused part of the code. The code in the fix should deal exclusively with the case that is not properly handled by ops.reshape. This is when the batch dimension does not have a static shape (shape == None) but all other dimensions have static shape. In that case we can determine the static shape of the dimension containing -1 form all dimensions that are not the batch dimension.

* applied changes according to pre-commit hook

* Added reshape_test test case that fails with original implementation and succeeds with fix.

* Added asserts for expected result.

* Fixed test name and added a further test with custom Model and Conv1D layer.

* applied changes according to pre-commit hook

* Fixed test to use a custom model and not a custom layer as indicated in the test name.

* Implemented suggested changes:

- Use original implementation from build method to avoid repeating the implementation in compute_reshape_output_shape.

- Set the self.built attribute to True in __init__ because no further build is required.

Additional change:

- Mention the optional use of -1 for a single dimension of the target_shape.

* Implemented suggested changes:

- removed explicit call to layer.build from the tests,
- changed the new test to use Model.compile and Model.fit to cover the corresponding API in the tests.

* Remove unused variable.

* Fixed line lengths in doc string.

* Marked test which uses fit method to require a trainable backend.

* Docs:

- Adapted according to review.
- Additionally, replaced "length" with "size" when referring to total number of elements.

* Fixed line length.

Author: roebel

keras/src/layers/reshaping/reshape.py

@@ -11,13 +11,12 @@ class Reshape(Layer):
 
     Args:
         target_shape: Target shape. Tuple of integers, does not include the
-            samples dimension (batch size).
+            samples dimension (batch size). One element of the `target_shape`
+            can be -1 in which case the missing value is inferred from the
+            size of the array and remaining dimensions.
 
     Input shape:
-        Arbitrary, although all dimensions in the input shape must be
-        known/fixed. Use the keyword argument `input_shape` (tuple of integers,
-        does not include the samples/batch size axis) when using this layer as
-        the first layer in a model.
+        Arbitrary, but required to be compatible with `target_shape`.
 
     Output shape:
         `(batch_size, *target_shape)`
@@ -29,15 +28,23 @@ class Reshape(Layer):
     >>> y.shape
     (None, 3, 4)
 
-    >>> # also supports shape inference using `-1` as dimension
+    >>> # another example with shape inference using `-1` as dimension
     >>> y = keras.layers.Reshape((-1, 2, 2))(x)
     >>> y.shape
     (None, 3, 2, 2)
     """
 
     def __init__(self, target_shape, **kwargs):
         super().__init__(**kwargs)
-        self.target_shape = tuple(target_shape)
+        target_shape = tuple(target_shape)
+        # test validity of target_shape
+        if target_shape.count(-1) > 1:
+            raise ValueError(
+                "The `target_shape` argument must not contain more than one "
+                f"`-1` value. Received: target_shape={target_shape}"
+            )
+        self.target_shape = target_shape
+        self.built = True
 
     def compute_output_shape(self, input_shape):
         return (
@@ -53,17 +60,17 @@ def compute_output_spec(self, inputs):
             shape=output_shape, dtype=inputs.dtype, sparse=inputs.sparse
         )
 
-    def build(self, input_shape):
-        sample_output_shape = operation_utils.compute_reshape_output_shape(
-            input_shape[1:], self.target_shape, "target_shape"
+    def call(self, inputs):
+        potentially_resolved_target_shape = (
+            operation_utils.compute_reshape_output_shape(
+                tuple(inputs.shape)[1:], self.target_shape, "target_shape"
+            )
         )
-        self._resolved_target_shape = tuple(
-            -1 if d is None else d for d in sample_output_shape
+        potentially_resolved_target_shape = tuple(
+            -1 if d is None else d for d in potentially_resolved_target_shape
         )
-
-    def call(self, inputs):
         return ops.reshape(
-            inputs, (ops.shape(inputs)[0],) + self._resolved_target_shape
+            inputs, (ops.shape(inputs)[0],) + potentially_resolved_target_shape
         )
 
     def get_config(self):

keras/src/layers/reshaping/reshape_test.py

@@ -1,8 +1,10 @@
 import pytest
 from absl.testing import parameterized
 
+from keras.src import Sequential
 from keras.src import backend
 from keras.src import layers
+from keras.src import ops
 from keras.src import testing
 from keras.src.backend.common.keras_tensor import KerasTensor
 
@@ -96,14 +98,19 @@ def test_reshape_with_dynamic_batch_size(self):
     def test_reshape_with_dynamic_batch_size_and_minus_one(self):
         input = KerasTensor((None, 6, 4))
         layer = layers.Reshape((-1, 8))
-        layer.build(input.shape)
         reshaped = backend.compute_output_spec(layer.__call__, input)
         self.assertEqual(reshaped.shape, (None, 3, 8))
 
+    def test_reshape_layer_with_varying_input_size_and_minus_one(self):
+        layer = layers.Reshape((-1, 8))
+        res = layer(ops.ones((1, 6, 4), dtype="float32"))
+        self.assertEqual(res.shape, (1, 3, 8))
+        res = layer(ops.ones((1, 10, 4), dtype="float32"))
+        self.assertEqual(res.shape, (1, 5, 8))
+
     def test_reshape_with_dynamic_dim_and_minus_one(self):
         input = KerasTensor((4, 6, None, 3))
         layer = layers.Reshape((-1, 3))
-        layer.build(input.shape)
         reshaped = backend.compute_output_spec(layer.__call__, input)
         self.assertEqual(reshaped.shape, (4, None, 3))
 
@@ -112,3 +119,20 @@ def test_reshape_sets_static_shape(self):
         reshaped = layers.Reshape((3, 5))(input_layer)
         # Also make sure the batch dim is not lost after reshape.
         self.assertEqual(reshaped.shape, (2, 3, 5))
+
+    @pytest.mark.requires_trainable_backend
+    def test_reshape_model_fit_with_varying_input_size_and_minus_one(self):
+        def generator():
+            yield (
+                ops.ones((1, 12, 2), dtype="float32"),
+                ops.zeros((1, 3, 8), dtype="float32"),
+            )
+            yield (
+                ops.ones((1, 20, 2), dtype="float32"),
+                ops.zeros((1, 5, 8), dtype="float32"),
+            )
+
+        layer = layers.Reshape((-1, 8))
+        model = Sequential([layer])
+        model.compile(loss="mean_squared_error")
+        model.fit(generator(), steps_per_epoch=2, epochs=1)