Don't lock functional models (#823)

* Don't lock functional models

* Add unit tests

keras_core/models/functional.py

@@ -161,6 +161,12 @@ def __init__(self, inputs, outputs, name=None, **kwargs):
         output_layers = [x._keras_history[0] for x in self.outputs]
         self.output_names = [x.name for x in output_layers]
 
+    def _lock_state(self):
+        # Unlike other layers, we allow Functional state to be mutable after
+        # build. E.g. to attach a layer to a model that is not part of the
+        # functional DAG.
+        pass
+
     @property
     def layers(self):
         layers = []

keras_core/models/functional_test.py

@@ -51,6 +51,16 @@ def test_scalar_input(self):
         out_val = model(in_val)
         self.assertAllClose(out_val, np.ones((2, 3)))
 
+    @pytest.mark.requires_trainable_backend
+    def test_mutable_state(self):
+        inputs = Input(shape=(3,), batch_size=2, name="input")
+        x = layers.Dense(5)(inputs)
+        outputs = layers.Dense(5)(x)
+        model = Functional(inputs, outputs)
+        # Allow attaching state to a model that isn't directly part of the DAG.
+        # Most useful for functional subclasses.
+        model.extra_layer = layers.Dense(5)
+
     @pytest.mark.requires_trainable_backend
     def test_basic_flow_multi_output(self):
         inputs = Input(shape=(3,), batch_size=2, name="input")