Version bump 3.11.2 and nnx fix #21565 (#21570)

* 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

* Version bump to 3.11.2

---------

Co-authored-by: Divyashree Sreepathihalli <[email protected]>

Author: laxmareddyp

.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"""
                 )
             )

keras/src/version.py

@@ -1,7 +1,7 @@
 from keras.src.api_export import keras_export
 
 # Unique source of truth for the version number.
-__version__ = "3.11.1"
+__version__ = "3.11.2"
 
 
 @keras_export("keras.version")