Skip to content

Version bump 3.11.2 and nnx fix #21565 #21570

New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

Sign up for GitHub

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Merged
merged 2 commits into from
Aug 11, 2025
Merged

Version bump 3.11.2 and nnx fix #21565 #21570

Show file tree
Hide file tree
Changes from all commits
Commits
Show all changes
2 commits
Select commit Hold shift + click to select a range
e8c64a7
Fix nnx object state (#21565)
divyashreepathihalli Aug 9, 2025
5fdfe07
Version bump to 3.11.2
laxmareddyp Aug 11, 2025
File filter

Filter by extension

Filter by extension
Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
2 changes: 1 addition & 1 deletion .github/workflows/actions.yml
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down
11 changes: 7 additions & 4 deletions keras/src/ops/operation.py
View file
Open in desktop
Original file line number Diff line number Diff line change
Expand Up @@ -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()
Comment on lines +126 to +129
Copy link
Contributor

Choose a reason for hiding this comment

The reason will be displayed to describe this comment to others. Learn more.

medium

The use of a broad try...except AttributeError block can be brittle. It might catch an AttributeError from an unexpected source within the PytreeState() call, potentially masking other issues. A more explicit check for the existence of nnx.pytreelib.PytreeState using hasattr would be more robust and clearer about the intent, which is to handle different versions of the flax library.

Suggested change
try:
vars(instance)["_pytree__state"] = nnx.pytreelib.PytreeState()
except AttributeError:
vars(instance)["_object__state"] = nnx.object.ObjectState()
if hasattr(nnx, "pytreelib") and hasattr(nnx.pytreelib, "PytreeState"):
vars(instance)["_pytree__state"] = nnx.pytreelib.PytreeState()
else:
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)))
Expand Down Expand Up @@ -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"""
)
)
Expand Down
2 changes: 1 addition & 1 deletion keras/src/version.py
View file
Open in desktop
Original file line number Diff line number Diff line change
@@ -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")
Expand Down
Loading