Repository: keras-team/keras (python)

Issue:
TimeDistributed layer with nested model no longer working in TensorFlow 2.16.1
With TensorFlow `2.15.1`, the following code works fine:

```python3
import numpy as np
from tensorflow.keras.layers import Input, TimeDistributed, Flatten
from tensorflow.keras.models import Model, Sequential

inputs = [Input((17, 4))]

nested_model = Sequential()
nested_model.add(Flatten(input_shape=(4,)))
nested_model.compile(optimizer='rmsprop', loss='categorical_crossentropy')
outputs = [TimeDistributed(nested_model)(inputs[0])]

model = Model(inputs=inputs, outputs=outputs)
model.compile(loss='mse', optimizer='adam')

x = np.random.random((1, 17, 4))
model.predict(x)
```

But with `2.16.1`, it no longer does:

```
NotImplementedError: Exception encountered when calling TimeDistributed.call().

Layer Sequential should implement `def compute_output_shape(self, input_shape)`.
```

Pull request description:
Enable `Sequential` and `Functional` model to be used as a submodel
Fix #19792

As titled.

To enable this, we can simply add `compute_output_shape` to `Sequential` and `Functional`.
The added unit test verifies this.

(I also encountered this issue when building the project https://github.com/james77777778/keras-image-models)

Patch:
diff --git a/.github/workflows/scripts/auto-assignment.js b/.github/workflows/scripts/auto-assignment.js
index 4b16ef1ea..89398a373 100644
--- a/.github/workflows/scripts/auto-assignment.js
+++ b/.github/workflows/scripts/auto-assignment.js
@@ -29,7 +29,7 @@ module.exports = async ({ github, context }) => {
   // Is this an issue? If so, assign the issue number. Otherwise, assign the PR number.
   if (context.payload.issue) {
     //assignee List for issues. 
-    assigneesList = ["SuryanarayanaY", "sachinprasadhs"];
+    assigneesList = ["mehtamansi29", "sachinprasadhs"];
     issueNumber = context.payload.issue.number;
   } else {
     //assignee List for PRs. 
diff --git a/keras/src/layers/reshaping/permute.py b/keras/src/layers/reshaping/permute.py
index fce4eb2d3..86580dfa0 100644
--- a/keras/src/layers/reshaping/permute.py
+++ b/keras/src/layers/reshaping/permute.py
@@ -14,7 +14,7 @@ class Permute(Layer):
     Args:
         dims: Tuple of integers. Permutation pattern does not include the
             batch dimension. Indexing starts at 1.
-            For instance, `(2, 1)` permutes the first and second dimensions
+            For instance, `(1, 3, 2)` permutes the second and third dimensions
             of the input.
 
     Input shape:
diff --git a/keras/src/models/functional.py b/keras/src/models/functional.py
index 85533d0a3..51d7eb021 100644
--- a/keras/src/models/functional.py
+++ b/keras/src/models/functional.py
@@ -181,6 +181,10 @@ class Functional(Function, Model):
         # From Function
         return super().compute_output_spec(inputs)
 
+    def compute_output_shape(self, input_shape):
+        # From Function
+        return super().compute_output_shape(input_shape)
+
     def build(self, input_shape):
         self.built = True
 
diff --git a/keras/src/models/sequential.py b/keras/src/models/sequential.py
index 9d7b8149e..821b12d20 100644
--- a/keras/src/models/sequential.py
+++ b/keras/src/models/sequential.py
@@ -252,6 +252,15 @@ class Sequential(Model):
             inputs = outputs
         return outputs
 
+    def compute_output_shape(self, input_shape):
+        if self._functional:
+            return self._functional.compute_output_shape(input_shape)
+        # Direct application
+        for layer in self.layers:
+            output_shape = layer.compute_output_shape(input_shape)
+            input_shape = output_shape
+        return output_shape
+
     @property
     def input_shape(self):
         if self._functional:
diff --git a/keras/src/ops/function.py b/keras/src/ops/function.py
index 49df4073e..f34ea3617 100644
--- a/keras/src/ops/function.py
+++ b/keras/src/ops/function.py
@@ -118,6 +118,20 @@ class Function(Operation):
             inputs, operation_fn=lambda op: op.compute_output_spec
         )
 
+    def compute_output_shape(self, input_shape):
+        # Wrap `input_shape` into the structure of KerasTensor to utilize
+        # `compute_output_spec`.
+        input_shape_struct = tree.map_shape_structure(
+            lambda x: KerasTensor(shape=x), input_shape
+        )
+        # Ensure that dtype and sparse settings are the same as self._inputs,
+        # because we only care about the shape in this function.
+        for x, x_ref in zip(tree.flatten(input_shape_struct), self._inputs):
+            x.dtype = x_ref.dtype
+            x.sparse = x_ref.sparse
+        output_spec = self.compute_output_spec(input_shape_struct)
+        return tree.map_structure(lambda x: x.shape, output_spec)
+
     def call(self, inputs):
         """Computes output tensors for new inputs."""
         self._assert_input_compatibility(inputs)
diff --git a/keras/src/utils/backend_utils.py b/keras/src/utils/backend_utils.py
index 4b34add28..3e974aa3e 100644
--- a/keras/src/utils/backend_utils.py
+++ b/keras/src/utils/backend_utils.py
@@ -60,6 +60,11 @@ class DynamicBackend:
         self._backend = backend or backend_module.backend()
 
     def set_backend(self, backend):
+        if backend not in ("tensorflow", "jax", "torch", "numpy"):
+            raise ValueError(
+                "Available backends are ('tensorflow', 'jax', 'torch' and "
+                f"'numpy'). Received: backend={backend}"
+            )
         self._backend = backend
 
     def reset(self):
@@ -71,25 +76,22 @@ class DynamicBackend:
 
     def __getattr__(self, name):
         if self._backend == "tensorflow":
-            from keras.src.backend import tensorflow as tf_backend
-
-            return getattr(tf_backend, name)
+            module = importlib.import_module("keras.src.backend.tensorflow")
+            return getattr(module, name)
         if self._backend == "jax":
-            from keras.src.backend import jax as jax_backend
-
-            return getattr(jax_backend, name)
+            module = importlib.import_module("keras.src.backend.jax")
+            return getattr(module, name)
         if self._backend == "torch":
-            from keras.src.backend import torch as torch_backend
-
-            return getattr(torch_backend, name)
+            module = importlib.import_module("keras.src.backend.torch")
+            return getattr(module, name)
         if self._backend == "numpy":
-            # TODO (ariG23498):
-            # The import `from keras.src.backend import numpy as numpy_backend`
-            # is not working. This is a temporary fix.
-            # The import is redirected to `keras.backend.numpy.numpy.py`
-            from keras.src import backend as numpy_backend
-
-            return getattr(numpy_backend, name)
+            if backend_module.backend() == "numpy":
+                return getattr(backend_module, name)
+            else:
+                raise NotImplementedError(
+                    "Currently, we cannot dynamically import the numpy backend "
+                    "because it would disrupt the namespace of the import."
+                )
 
 
 @keras_export("keras.config.set_backend")


Test patch:
diff --git a/keras/src/models/functional_test.py b/keras/src/models/functional_test.py
index 7d17abee6..de425d445 100644
--- a/keras/src/models/functional_test.py
+++ b/keras/src/models/functional_test.py
@@ -118,6 +118,20 @@ class FunctionalTest(testing.TestCase):
         out_val = model(in_val)
         self.assertEqual(out_val.shape, (2, 4))
 
+    def test_basic_flow_as_a_submodel(self):
+        # Build submodel
+        submodel_inputs = Input([4])
+        submodel_outputs = layers.Flatten()(submodel_inputs)
+        submodel = Model(submodel_inputs, submodel_outputs)
+
+        inputs = Input((None, 4))
+        outputs = layers.TimeDistributed(submodel)(inputs)
+        model = Model(inputs=inputs, outputs=outputs)
+
+        x = np.random.random((2, 3, 4))
+        y = model(x)
+        self.assertEqual(y.shape, (2, 3, 4))
+
     @pytest.mark.requires_trainable_backend
     def test_named_input_dict_io(self):
         input_a = Input(shape=(3,), batch_size=2, name="a")
diff --git a/keras/src/models/sequential_test.py b/keras/src/models/sequential_test.py
index 6ce5ff5f2..ce51d7238 100644
--- a/keras/src/models/sequential_test.py
+++ b/keras/src/models/sequential_test.py
@@ -8,6 +8,7 @@ from keras.src import layers
 from keras.src import testing
 from keras.src.layers.core.input_layer import Input
 from keras.src.models.functional import Functional
+from keras.src.models.model import Model
 from keras.src.models.sequential import Sequential
 
 
@@ -135,6 +136,20 @@ class SequentialTest(testing.TestCase):
         y = model(x)
         self.assertEqual(y.shape, (3, 4))
 
+    def test_basic_flow_as_a_submodel(self):
+        # Build submodel
+        submodel = Sequential()
+        submodel.add(layers.Flatten())
+        self.assertFalse(submodel.built)
+
+        inputs = Input((None, 4))
+        outputs = layers.TimeDistributed(submodel)(inputs)
+        model = Model(inputs=inputs, outputs=outputs)
+
+        x = np.random.random((2, 3, 4))
+        y = model(x)
+        self.assertEqual(y.shape, (2, 3, 4))
+
     def test_dict_inputs(self):
         class DictLayer(layers.Layer):
             def call(self, inputs):
@@ -271,3 +286,8 @@ class SequentialTest(testing.TestCase):
             ValueError, "can only have a single positional"
         ):
             model.build((None, 2))
+
+    def test_compute_output_shape(self):
+        layer = Sequential([layers.Dense(4), layers.Dense(8)])
+        output_shape = layer.compute_output_shape((1, 2))
+        self.assertEqual(output_shape, (1, 8))
diff --git a/keras/src/ops/function_test.py b/keras/src/ops/function_test.py
index d1ebc838d..54160fa8f 100644
--- a/keras/src/ops/function_test.py
+++ b/keras/src/ops/function_test.py
@@ -55,6 +55,11 @@ class FunctionTest(testing.TestCase):
         self.assertIsInstance(out, keras_tensor.KerasTensor)
         self.assertEqual(out.shape, (4, 3))
 
+        # Test with compute_output_shape
+        out = fn.compute_output_shape((None, 3))
+        self.assertIsInstance(out, tuple)
+        self.assertEqual(out, (None, 3))
+
         # Test with call
         out = fn(keras_tensor.KerasTensor((4, 3)))
         self.assertIsInstance(out, keras_tensor.KerasTensor)
diff --git a/keras/src/utils/backend_utils_test.py b/keras/src/utils/backend_utils_test.py
new file mode 100644
index 000000000..ef5c8dfa2
--- /dev/null
+++ b/keras/src/utils/backend_utils_test.py
@@ -0,0 +1,53 @@
+import numpy as np
+from absl.testing import parameterized
+
+from keras.src import backend
+from keras.src import testing
+from keras.src.utils import backend_utils
+
+
+class BackendUtilsTest(testing.TestCase, parameterized.TestCase):
+    @parameterized.named_parameters(
+        ("numpy", "numpy"),
+        ("jax", "jax"),
+        ("tensorflow", "tensorflow"),
+        ("torch", "torch"),
+    )
+    def test_dynamic_backend(self, name):
+        dynamic_backend = backend_utils.DynamicBackend()
+        x = np.random.uniform(size=[1, 2, 3])
+
+        if name == "numpy":
+            dynamic_backend.set_backend(name)
+            if backend.backend() != "numpy":
+                with self.assertRaisesRegex(
+                    NotImplementedError,
+                    "Currently, we cannot dynamically import the numpy backend",
+                ):
+                    y = dynamic_backend.numpy.log10(x)
+            else:
+                y = dynamic_backend.numpy.log10(x)
+                self.assertIsInstance(y, np.ndarray)
+        elif name == "jax":
+            import jax
+
+            dynamic_backend.set_backend(name)
+            y = dynamic_backend.numpy.log10(x)
+            self.assertIsInstance(y, jax.Array)
+        elif name == "tensorflow":
+            import tensorflow as tf
+
+            dynamic_backend.set_backend(name)
+            y = dynamic_backend.numpy.log10(x)
+            self.assertIsInstance(y, tf.Tensor)
+        elif name == "torch":
+            import torch
+
+            dynamic_backend.set_backend(name)
+            y = dynamic_backend.numpy.log10(x)
+            self.assertIsInstance(y, torch.Tensor)
+
+    def test_dynamic_backend_invalid_name(self):
+        dynamic_backend = backend_utils.DynamicBackend()
+        with self.assertRaisesRegex(ValueError, "Available backends are"):
+            dynamic_backend.set_backend("abc")