Fix NNX tests #21884
Fix NNX tests #21884
Conversation
Summary of ChangesHello @divyashreepathihalli, I'm Gemini Code Assist1! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed! This pull request addresses critical test failures in the Keras NNX integration tests by updating core dependencies, Flax and JAX, to their latest versions. The primary fix involves adapting to changes in Flax's Highlights
Using Gemini Code AssistThe full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips. Invoking Gemini You can request assistance from Gemini at any point by creating a comment using either
Customization To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a Limitations & Feedback Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here. You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension. Footnotes
|
![gemini-code-assist[bot]](https://avatars.githubusercontent.com/in/956858?s=60&v=4){kind=small}
There was a problem hiding this comment.
Code Review
This pull request effectively addresses test failures in NNX integration by updating JAX and Flax dependencies and registering NnxVariable as a custom JAX PyTree node. This ensures Keras-specific attributes are preserved during NNX operations. The changes are well-reasoned. I have two suggestions for improvement: one to update an outdated comment in requirements.txt, and another to enhance the robustness of the new __setattr__ fallback mechanism in keras/src/backend/jax/core.py to prevent potential state inconsistencies.
| jax[cpu]==0.8.1 | ||
| flax==0.12.1 |
There was a problem hiding this comment.
With the JAX version update, the comment on line 14 explaining the pin to 0.5.0 is now outdated. Please update or remove it to avoid confusion for future developers.
Codecov Report❌ Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #21884 +/- ##
===========================================
- Coverage 82.57% 61.46% -21.11%
===========================================
Files 577 578 +1
Lines 59599 59813 +214
Branches 9351 9387 +36
===========================================
- Hits 49213 36767 -12446
- Misses 7978 20720 +12742
+ Partials 2408 2326 -82
Flags with carried forward coverage won't be shown. Click here to find out more. ☔ View full report in Codecov by Sentry. 🚀 New features to boost your workflow:
|
hertschuh
left a comment
hertschuh
left a comment
There was a problem hiding this comment.
Thanks for the PR!
The Flax tests are failing though
| # Mirror Keras attributes to _var_metadata to ensure persistence | ||
| # if the Pytree registration is not respected by NNX. | ||
| if ( | ||
| name != "_var_metadata" | ||
| and name not in ("_raw_value", "_trace_state") | ||
| and hasattr(self, "_var_metadata") | ||
| ): | ||
| self._var_metadata[name] = value | ||
|
|
||
| object.__setattr__(self, name, value) | ||
|
|
||
| NnxVariable.__setattr__ = __setattr__ |
There was a problem hiding this comment.
Why do you do NnxVariable.__setattr__ = __setattr__ instead of simply declaring def __setattr__(self, name, value): directly within the NnxVariable class?
There was a problem hiding this comment.
This is needed to mirror Keras attributes into _var_metadata, so that they persist during NNX's Pytree registration, which might otherwise be missed if defined directly within the class due to the complex multiple inheritance
| # Note that we test against the latest JAX on GPU. | ||
| jax[cpu]==0.5.0 | ||
| flax | ||
| jax[cpu]==0.8.1 |
There was a problem hiding this comment.
This is promising, but the TensorFlow CPU tests are failing. Upgrade TF version too?
divyashreepathihalli
force-pushed
the
nnx_fix
branch
from
80cf2b3 to
4143934
Compare
google-ml-butler
bot
added
kokoro:force-run
ready to pull
google-ml-butler
bot
removed
the
ready to pull
divyashreepathihalli
force-pushed
the
nnx_fix
branch
from
9fedd6d to
10ba8ec
Compare
divyashreepathihalli
merged commit 9acaaf5
into
keras-team:master
5 participants
This PR updates the Flax dependency to 0.12.1 and JAX to 0.8.1.
It addresses test failures in the NNX integration tests (
keras/src/backend/jax/core_test.py) caused by changes in Flax'snnx.Variableimplementation.nnx.Variablenow relies heavily on slots and metadata, causing Keras-specific attributes stored in__dict__to be lost duringnnx.split/merge.Fix: Registering
NnxVariableas a custom Pytree node with JAX. This ensures that the instance's__dict__(excluding the raw value) is treated as auxiliary data (static metadata) and preserved in the GraphDef.