fix for keras-team#21120

Author: martin-gorner

keras/src/layers/core/composite_layer_test.py

@@ -1198,6 +1198,31 @@ def replace_fn(layer, *args, **kwargs):
 
         self.assertAllClose(model(data), model2(data))
 
+    def test_clone_recursive(self):
+        def layer_fn1(inputs):
+            return layers.Dense(32, name="dense_2")(inputs)
+        
+        layer1 = CompositeLayer(layer_fn1, name="sub")
+
+        def layer_fn2(inputs):
+            return layer1(inputs)
+
+        slayer = CompositeLayer(layer_fn2, name="subfunc")
+
+        inputs = layers.Input(shape=(16, 32))
+        outputs = slayer(inputs)
+        model = Model(inputs, outputs)
+
+        def call_fn(layer, *args, **kwargs):
+            if isinstance(layer, layers.Dense):
+                new_layer = layers.Dense(layer.units, name="dense_modified")
+                return new_layer(*args, **kwargs)
+            return layer(*args, **kwargs)
+
+        new_model = clone_model(model, call_function=call_fn, recursive=True)
+        sub = new_model.get_layer("subfunc").get_layer("sub")
+        self.assertEqual(sub.layers[1].name, "dense_modified")
+        
     def test_build_twice(self):
         def layer_fn(inputs):
             return layers.Dense(5)(inputs)

keras/src/models/cloning.py

@@ -276,6 +276,7 @@ def wrapped_clone_function(layer):
                     clone_function=clone_function,
                     call_function=call_function,
                     cache=cache,
+                    recursive=True
                 )
                 cache[id(layer)] = clone
                 return clone

keras/src/models/cloning_test.py

@@ -279,6 +279,29 @@ def clone_function(layer):
                 self.assertFalse(hasattr(l1, "flag"))
                 self.assertTrue(hasattr(l2, "flag"))
 
+    def test_recursive_level_2(self):
+        inputs = layers.Input(shape=(16, 32))
+        outputs = layers.Dense(32, name="dense_2")(inputs)
+        layer1 = models.Model(inputs, outputs, name="sub")
+
+        inputs = layers.Input(shape=(16, 32))
+        outputs = layer1(inputs)
+        slayer = models.Model(inputs, outputs, name="subfunc")
+
+        inputs = layers.Input(shape=(16, 32))
+        outputs = slayer(inputs)
+        model = models.Model(inputs, outputs)
+
+        def call_fn(layer, *args, **kwargs):
+            if isinstance(layer, layers.Dense):
+                new_layer = layers.Dense(layer.units, name="dense_modified")
+                return new_layer(*args, **kwargs)
+            return layer(*args, **kwargs)
+
+        new_model = clone_model(model, call_function=call_fn, recursive=True)
+        sub = new_model.get_layer("subfunc").get_layer("sub")
+        self.assertEqual(sub.layers[1].name, "dense_modified")
+
     def test_compiled_model_cloning(self):
         model = models.Sequential()
         model.add(layers.Input((3,)))