add object state in functional

Author: divyashreepathihalli

keras/src/models/functional.py

@@ -98,8 +98,13 @@ class Functional(Function, Model):
     """
 
     def __new__(cls, *args, **kwargs):
-        obj = super().__new__(cls, *args, **kwargs)
-        return obj
+        
+        if backend.backend() == "jax" and is_nnx_enabled():
+            instance = super(Functional, cls).__new__(cls)
+            from flax import nnx
+
+            vars(instance)["_object__state"] = nnx.object.ObjectState()
+        return typing.cast(cls, super().__new__(cls))
 
 
     @tracking.no_automatic_dependency_tracking