Repository: keras-team/keras

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


Discussion:
Are you satisfied with the resolution of your issue?
<a href="https://docs.google.com/forms/d/e/1FAIpQLSdHag0RVFS7UXzZkKcsFCKOcX8raCupKK9RHSlYxp5U8lSJbQ/viewform?entry.492125872=Yes&entry.243948740=https://github.com/keras-team/keras/issues/19913">Yes</a>
<a href="https://docs.google.com/forms/d/e/1FAIpQLSdHag0RVFS7UXzZkKcsFCKOcX8raCupKK9RHSlYxp5U8lSJbQ/viewform?entry.492125872=No&entry.243948740=https://github.com/keras-team/keras/issues/19913">No</a>




Patch:
diff --git a/keras/src/export/export_lib.py b/keras/src/export/export_lib.py
index 02714c55c0ad..be75c06cd3d7 100644
--- a/keras/src/export/export_lib.py
+++ b/keras/src/export/export_lib.py
@@ -654,13 +654,18 @@ def make_tensor_spec(structure):
         # 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}"


Test patch:
diff --git a/keras/src/export/export_lib_test.py b/keras/src/export/export_lib_test.py
index 29504cfb2b1b..7b4b7d332dcd 100644
--- a/keras/src/export/export_lib_test.py
+++ b/keras/src/export/export_lib_test.py
@@ -196,6 +196,22 @@ def call(self, inputs):
         )
         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):