Fix issue with deserialization of nested and shared Functional models.

Author: fchollet

keras/models/functional.py

@@ -480,7 +480,7 @@ def process_node(layer, node_data):
         layer(*args, **kwargs)
 
     def process_layer(layer_data):
-        """Deserializes a layer, then call it on appropriate inputs.
+        """Deserializes a layer and index its inbound nodes.
 
         Args:
             layer_data: layer config dict.
@@ -553,6 +553,9 @@ def process_layer(layer_data):
     def get_tensor(layer_name, node_index, tensor_index):
         assert layer_name in created_layers
         layer = created_layers[layer_name]
+        if isinstance(layer, Functional):
+            # Functional models start out with a built-in node.
+            node_index -= 1
         layer_output_tensors = layer._inbound_nodes[node_index].output_tensors
         return layer_output_tensors[tensor_index]
 
@@ -613,8 +616,9 @@ def serialize_keras_tensor(x):
         if isinstance(x, backend.KerasTensor):
             operation, node_index, tensor_index = x._keras_history
             irrelevant_node_count = 0
-            for node in operation._inbound_nodes[:node_index]:
-                if node not in own_nodes:
+            for i, node in enumerate(operation._inbound_nodes[:node_index]):
+                node_key = make_node_key(operation, i)
+                if node_key not in own_nodes:
                     irrelevant_node_count += 1
             x._keras_history = KerasHistory(
                 operation, node_index - irrelevant_node_count, tensor_index

keras/ops/node.py

@@ -82,7 +82,7 @@ def __init__(
         self.is_input = not self.arguments.keras_tensors
 
     def __repr__(self):
-        return f"<Node operation={self.operation}, id={id(self)}>"
+        return f"<Node operation={self.operation.name}, id={id(self)}>"
 
     @property
     def input_tensors(self):

keras/saving/saving_lib_test.py

@@ -861,3 +861,48 @@ def test_legacy_h5_format(self):
         new_model = keras.saving.load_model(temp_filepath)
         out = new_model(x)
         self.assertAllClose(ref_out, out, atol=1e-6)
+
+    def test_nested_functional_model_saving(self):
+        def func(in_size=4, out_size=2, name=None):
+            inputs = keras.layers.Input(shape=(in_size,))
+            outputs = keras.layers.Dense(out_size)((inputs))
+            return keras.Model(inputs, outputs=outputs, name=name)
+
+        input_a, input_b = keras.Input((4,)), keras.Input((4,))
+        out_a = func(out_size=2, name="func_a")(input_a)
+        out_b = func(out_size=3, name="func_b")(input_b)
+        model = keras.Model([input_a, input_b], outputs=[out_a, out_b])
+
+        temp_filepath = os.path.join(self.get_temp_dir(), "nested_func.keras")
+        model.save(temp_filepath)
+        new_model = keras.saving.load_model(temp_filepath)
+        x = [np.random.random((2, 4))], np.random.random((2, 4))
+        ref_out = model(x)
+        out = new_model(x)
+        self.assertAllClose(ref_out[0], out[0])
+        self.assertAllClose(ref_out[1], out[1])
+
+    def test_nested_shared_functional_model_saving(self):
+        def func(in_size=4, out_size=2, name=None):
+            inputs = keras.layers.Input(shape=(in_size,))
+            outputs = keras.layers.Dense(out_size)((inputs))
+            return keras.Model(inputs, outputs=outputs, name=name)
+
+        inputs = [keras.Input((4,)), keras.Input((4,))]
+        func_shared = func(out_size=4, name="func_shared")
+        shared_a = func_shared(inputs[0])
+        shared_b = func_shared(inputs[1])
+        out_a = keras.layers.Dense(2)(shared_a)
+        out_b = keras.layers.Dense(2)(shared_b)
+        model = keras.Model(inputs, outputs=[out_a, out_b])
+
+        temp_filepath = os.path.join(
+            self.get_temp_dir(), "nested_shared_func.keras"
+        )
+        model.save(temp_filepath)
+        new_model = keras.saving.load_model(temp_filepath)
+        x = [np.random.random((2, 4))], np.random.random((2, 4))
+        ref_out = model(x)
+        out = new_model(x)
+        self.assertAllClose(ref_out[0], out[0])
+        self.assertAllClose(ref_out[1], out[1])