Repository: keras-team/keras (Python)
Task: Bug Fix (2024-06-25 14:03:04+00:00)

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:

Patch:
diff --git a/keras/src/export/export_lib.py b/keras/src/export/export_lib.py
--- 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
--- 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):