Fix nnx object state (#21565)

* Update operation.py

* Update actions.yml

* Update operation.py

* Update actions.yml

* Update operation.py

* Update operation.py

* Update operation.py

* fix test

* code reformat

Author: divyashreepathihalli

.github/workflows/actions.yml

@@ -57,7 +57,7 @@ jobs:
         run: |
           pip install -r requirements.txt --progress-bar off --upgrade
           if [ "${{ matrix.nnx_enabled }}" == "true" ]; then
-            pip install --upgrade flax>=0.11.0
+            pip install --upgrade flax>=0.11.1
           fi
           pip uninstall -y keras keras-nightly
           pip install -e "." --progress-bar off --upgrade

keras/src/ops/operation.py

@@ -123,7 +123,11 @@ def __new__(cls, *args, **kwargs):
         if backend.backend() == "jax" and is_nnx_enabled():
             from flax import nnx
 
-            vars(instance)["_object__state"] = nnx.object.ObjectState()
+            try:
+                vars(instance)["_pytree__state"] = nnx.pytreelib.PytreeState()
+            except AttributeError:
+                vars(instance)["_object__state"] = nnx.object.ObjectState()
+
         # Generate a config to be returned by default by `get_config()`.
         arg_names = inspect.getfullargspec(cls.__init__).args
         kwargs.update(dict(zip(arg_names[1 : len(args) + 1], args)))
@@ -206,10 +210,9 @@ def __init__(self, arg1, arg2, **kwargs):
 
             def get_config(self):
                 config = super().get_config()
-                config.update({{
-                    "arg1": self.arg1,
+                config.update({"arg1": self.arg1,
                     "arg2": self.arg2,
-                }})
+                })
                 return config"""
                 )
             )