fix legacy model saving & reloading with axis argument in its layer (#20973)

* fix legacy model saving & relaoding with axis arg in layer

* fix formatting issue

* add temp_file_path

Author: Surya2k1

keras/src/legacy/saving/legacy_h5_format_test.py

@@ -279,6 +279,21 @@ class RegisteredSubLayer(layers.Layer):
             self.assertIsInstance(loaded_layer.sublayers[1], RegisteredSubLayer)
             self.assertEqual(loaded_layer.sublayers[1].name, "MySubLayer")
 
+    def test_model_loading_with_axis_arg(self):
+        input1 = layers.Input(shape=(1, 4), name="input1")
+        input2 = layers.Input(shape=(1, 4), name="input2")
+        concat1 = layers.Concatenate(axis=1)([input1, input2])
+        output = layers.Dense(1, activation="sigmoid")(concat1)
+        model = models.Model(inputs=[input1, input2], outputs=output)
+        model.compile(
+            optimizer="adam", loss="binary_crossentropy", metrics=["accuracy"]
+        )
+        temp_filepath = os.path.join(
+            self.get_temp_dir(), "model_with_axis_arg.h5"
+        )
+        legacy_h5_format.save_model_to_hdf5(model, temp_filepath)
+        legacy_h5_format.load_model_from_hdf5(temp_filepath)
+
 
 @pytest.mark.requires_trainable_backend
 @pytest.mark.skipif(tf_keras is None, reason="Test requires tf_keras")

keras/src/legacy/saving/saving_utils.py

@@ -63,8 +63,11 @@ def model_from_config(config, custom_objects=None):
             config["config"]["input_shape"] = batch_input_shape
 
     axis = config["config"].pop("axis", None)
-    if axis is not None and isinstance(axis, list) and len(axis) == 1:
-        config["config"]["axis"] = int(axis[0])
+    if axis is not None:
+        if isinstance(axis, list) and len(axis) == 1:
+            config["config"]["axis"] = int(axis[0])
+        elif isinstance(axis, (int, float)):
+            config["config"]["axis"] = int(axis)
 
     # Handle backwards compatibility for Keras lambdas
     if config["class_name"] == "Lambda":