Keras <> NNX integration #21252
Keras <> NNX integration #21252
Conversation
Codecov ReportAttention: Patch coverage is
Additional details and impacted files@@ Coverage Diff @@
## master #21252 +/- ##
==========================================
- Coverage 82.67% 82.62% -0.06%
==========================================
Files 565 565
Lines 55064 55494 +430
Branches 8569 8691 +122
==========================================
+ Hits 45525 45851 +326
- Misses 7441 7528 +87
- Partials 2098 2115 +17
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:
|
keras/src/random/random_test.py
Outdated
| import jax.numpy as jnp | ||
|
|
||
| x = ops.ones(3) | ||
|
|
||
| @jax.jit | ||
| @nnx.jit |
There was a problem hiding this comment.
Would the integration prevent the use of jax.jit with Keras layers?
There was a problem hiding this comment.
yes! it would only work with nnx.jit for now ( They might be working on adding support for jax.jit)
There was a problem hiding this comment.
Added nnx as a opt in with this flag - os.environ["KERAS_NNX_ENABLED"]
keras/src/backend/config.py
Outdated
| @@ -230,6 +233,25 @@ def is_flash_attention_enabled(): | |||
| return global_state.get_global_attribute("flash_attention", default=None) | |||
|
|
|||
|
|
|||
| @keras_export("keras.config.is_nnx_backend_enabled") | |||
| def is_nnx_backend_enabled(): | |||
There was a problem hiding this comment.
Oh you renamed the setter set_nnx_enabled, but not the getter is_nnx_enabled.
| @@ -1533,7 +1539,19 @@ def __setattr__(self, name, value): | |||
| if not hasattr(self, "_tracker"): | |||
| self._initialize_tracker() | |||
| value = self._tracker.track(value) | |||
| return super().__setattr__(name, value) | |||
|
|
|||
| # NNX-specific bypass for `_called` and `built` attributes | |||
There was a problem hiding this comment.
Oh, I thought we were fixing that in NNX, is that not the case?
There was a problem hiding this comment.
Apparently not.
|
|
||
| vars(instance)["_object__state"] = nnx.object.ObjectState() |
There was a problem hiding this comment.
The comment got lost, can we move this to Layer?
There was a problem hiding this comment.
Unfortunately that did not work. It is strange. Linking my comment here
#21252 (comment)
divyashreepathihalli
force-pushed
the
nnx
branch
from
f652cb9 to
d544a0b
Compare
|
Noticed the following for when we use a custom NNX training loop.
Here is a notebook demonstrating the same: https://colab.research.google.com/drive/1hvGYN00aETHlrwjUlLCjfwVJrCcQqySN?resourcekey=0-oyfeCvzJzqOf0aLroY4rhg&usp=sharing. Let me know if I've messed up somewhere! Trying to figure out why it's happening. Will changing the |
The PR integrates NNX into JAX backend!
The following snippet shows how you would enable the nnx backend
Demo colab here : https://colab.sandbox.google.com/drive/1mK-4qbce2HGRIkcb4v5n4niWGDezL_6n#scrollTo=m-ZH9Mpnphfz
Added a github workflow action for nnx backend. Note this will fail - because this needs a new release of flax to work.