Repository: keras-team/keras (python)

Issue:
Unable to export reloaded model
Saving and reloading model makes it impossible to export it as a SavedModel artifact. 
Reloaded model has shapes defined as lists while export function expect tuples. 
Casting the shape to tuple in this particular place resolves the issue, but there may be other errors related to this in other places, though.

Steps to reproduce:

1) Make a subclassed model (maybe reproducible with Functional too?)
2) Save the model as `.keras`
3) Reload `.keras`
4) Try to `model.export()` on your reloaded model

Here's the notebook with the same steps for your convenience:
https://colab.research.google.com/drive/1oO4JxoYK4I4UO0VdyYPAlCQY9fT1pYlw

Pull request description:
Fix `export_lib.make_tensor_spec`
Fixes #19913 

a simple repro from the original issue with minor changes:

```
import keras
import numpy as np

class SomeModel(keras.Model):
    def __init__(self, output_channels: int):
        super().__init__()
        self.output_channels = output_channels
        self.fc1 = keras.layers.Dense(256, activation="relu")
        self.fc2 = keras.layers.Dense(
            self.output_channels * 2,
            bias_initializer="ones",
            kernel_initializer="zeros",
        )

    def build(self, input_shape):
        super().build(input_shape)
        self.fc1.build(input_shape)
        fc1_out = self.fc1.compute_output_shape(input_shape)
        self.fc2.build(fc1_out)

    def call(self, inputs):
        features = self.fc1(inputs)
        return self.fc2(features)


class SomeParentModel(keras.Model):
    def __init__(self, *args, **kwargs):
        super().__init__(*args, **kwargs)
        self.fx = SomeModel(20)

    def call(self, im_batch):
        features = self.fx(im_batch)
        ...
        return features

parent_inst = SomeParentModel()
random_data = np.random.random((1, 512)).astype('float32')

parent_inst(random_data)
parent_inst.save('test_model.keras')

reloaded = keras.models.load_model('test_model.keras', custom_objects={'SomeParentModel': SomeParentModel})
reloaded(random_data)
reloaded.export('test')
```




Patch:
diff --git a/keras/src/export/export_lib.py b/keras/src/export/export_lib.py
index 02714c55c..be75c06cd 100644
--- a/keras/src/export/export_lib.py
+++ b/keras/src/export/export_lib.py
@@ -654,13 +654,18 @@ def _get_input_signature(model):
         # into plain Python structures because they don't work with jax2tf/JAX.
         if isinstance(structure, dict):
             return {k: make_tensor_spec(v) for k, v in structure.items()}
-        if isinstance(structure, (list, tuple)):
+        elif isinstance(structure, tuple):
             if all(isinstance(d, (int, type(None))) for d in structure):
                 return tf.TensorSpec(
                     shape=(None,) + structure[1:], dtype=model.input_dtype
                 )
-            result = [make_tensor_spec(v) for v in structure]
-            return tuple(result) if isinstance(structure, tuple) else result
+            return tuple(make_tensor_spec(v) for v in structure)
+        elif isinstance(structure, list):
+            if all(isinstance(d, (int, type(None))) for d in structure):
+                return tf.TensorSpec(
+                    shape=[None] + structure[1:], dtype=model.input_dtype
+                )
+            return [make_tensor_spec(v) for v in structure]
         else:
             raise ValueError(
                 f"Unsupported type {type(structure)} for {structure}"
diff --git a/keras/src/layers/rnn/time_distributed.py b/keras/src/layers/rnn/time_distributed.py
index cac954a4f..e61274d96 100644
--- a/keras/src/layers/rnn/time_distributed.py
+++ b/keras/src/layers/rnn/time_distributed.py
@@ -72,8 +72,8 @@ class TimeDistributed(Wrapper):
         self.built = True
 
     def call(self, inputs, training=None, mask=None):
-        input_shape = inputs.shape
-        mask_shape = None if mask is None else tuple(mask.shape)
+        input_shape = ops.shape(inputs)
+        mask_shape = None if mask is None else ops.shape(mask)
         batch_size = input_shape[0]
         timesteps = input_shape[1]
 


Test patch:
diff --git a/keras/src/export/export_lib_test.py b/keras/src/export/export_lib_test.py
index 29504cfb2..7b4b7d332 100644
--- a/keras/src/export/export_lib_test.py
+++ b/keras/src/export/export_lib_test.py
@@ -196,6 +196,22 @@ class ExportArchiveTest(testing.TestCase, parameterized.TestCase):
         )
         revived_model.serve(bigger_input)
 
+        # Test with keras.saving_lib
+        temp_filepath = os.path.join(
+            self.get_temp_dir(), "exported_model.keras"
+        )
+        saving_lib.save_model(model, temp_filepath)
+        revived_model = saving_lib.load_model(
+            temp_filepath,
+            {
+                "TupleModel": TupleModel,
+                "ArrayModel": ArrayModel,
+                "DictModel": DictModel,
+            },
+        )
+        self.assertAllClose(ref_output, revived_model(ref_input))
+        export_lib.export_model(revived_model, self.get_temp_dir())
+
     def test_model_with_multiple_inputs(self):
 
         class TwoInputsModel(models.Model):
diff --git a/keras/src/ops/numpy_test.py b/keras/src/ops/numpy_test.py
index 76952cef3..2bfe4a637 100644
--- a/keras/src/ops/numpy_test.py
+++ b/keras/src/ops/numpy_test.py
@@ -2083,6 +2083,7 @@ class NumpyTwoInputOpsCorretnessTest(testing.TestCase, parameterized.TestCase):
         not backend.SUPPORTS_SPARSE_TENSORS,
         reason="Backend does not support sparse tensors.",
     )
+    @pytest.mark.skipif(testing.tensorflow_uses_gpu(), reason="Segfault")
     def test_matmul_sparse(self, dtype, x_shape, y_shape, x_sparse, y_sparse):
         if backend.backend() == "tensorflow":
             import tensorflow as tf