Keras <> NNX integration #21252

divyashreepathihalli · 2025-05-05T23:04:56Z

The PR integrates NNX into JAX backend!

The following snippet shows how you would enable the nnx backend

import os
os.environ["KERAS_BACKEND"]="jax"
os.environ["KERAS_NNX_ENABLED"]="true"
import keras

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.

codecov-commenter · 2025-05-05T23:10:19Z

Codecov Report

Attention: Patch coverage is 29.13907% with 107 lines in your changes missing coverage. Please review.

Project coverage is 82.72%. Comparing base (3554825) to head (0426eee).
Report is 1 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/src/backend/jax/core.py	6.49%	71 Missing and 1 partial ⚠️
keras/src/backend/config.py	52.38%	8 Missing and 2 partials ⚠️
keras/src/ops/function.py	50.00%	6 Missing and 2 partials ⚠️
keras/src/backend/jax/trainer.py	66.66%	2 Missing and 2 partials ⚠️
keras/src/layers/layer.py	50.00%	2 Missing and 2 partials ⚠️
keras/src/backend/jax/layer.py	50.00%	2 Missing and 1 partial ⚠️
keras/src/ops/operation.py	25.00%	2 Missing and 1 partial ⚠️
keras/src/backend/common/variables.py	0.00%	1 Missing and 1 partial ⚠️
keras/api/_tf_keras/keras/config/__init__.py	0.00%	1 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #21252      +/-   ##
==========================================
- Coverage   82.87%   82.72%   -0.16%     
==========================================
  Files         567      567              
  Lines       56073    56214     +141     
  Branches     8756     8786      +30     
==========================================
+ Hits        46470    46501      +31     
- Misses       7459     7556      +97     
- Partials     2144     2157      +13

Flag	Coverage Δ
keras	`82.52% <29.13%> (-0.16%)`	⬇️
keras-jax	`63.92% <29.13%> (-0.11%)`	⬇️
keras-numpy	`58.41% <21.19%> (-0.11%)`	⬇️
keras-openvino	`34.56% <19.86%> (-0.05%)`	⬇️
keras-tensorflow	`64.34% <21.85%> (-0.12%)`	⬇️
keras-torch	`63.97% <21.85%> (-0.12%)`	⬇️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

fchollet · 2025-05-10T22:46:02Z

keras/src/random/random_test.py

        x = ops.ones(3)

-        @jax.jit
+        @nnx.jit


Would the integration prevent the use of jax.jit with Keras layers?

yes! it would only work with nnx.jit for now ( They might be working on adding support for jax.jit)

Added nnx as a opt in with this flag - os.environ["KERAS_NNX_ENABLED"]

mattdangerw

Left some final comments on the config and configuration part. All good on my end once fixed.

.github/workflows/actions.yml

.github/workflows/config/nnx/keras.json

integration_tests/import_test.py

keras/src/backend/config.py

mattdangerw · 2025-07-24T20:47:20Z

keras/src/backend/config.py

+
+if "KERAS_NNX_ENABLED" in os.environ:
+    env_val = os.environ["KERAS_NNX_ENABLED"].lower()
+    if env_val:


This seems bad, don't we want to check if val.lower() = "true" or "1" here? Otherwise KERAS_NNX_ENABLED=False will enable nnx.

hertschuh · 2025-07-24T16:34:55Z

keras/src/backend/common/variables.py

        if self._value is not None:
+            # If NNX is enabled, it's possible the variable was already
+            # initialized by a concrete call. In this case,
+            # _deferred_initialize becomes a no-op for this variable.


But this comment is incorrect. It raises an error in this case, so it's not a no-op.

updated the comment and also if nnx is enabled, I added a return instead of retuning an error.

mattdangerw · 2025-07-25T01:06:51Z

keras/src/backend/config.py

+
+    # Apply basic configs that don't cause circular import
    set_floatx(_floatx)
+    _NNX_ENABLED = _nnx_enabled_config


can we literally call set_nnx_enabled(_nnx_enabled_config) (just going for consistency)

this leads to circular import error

from keras.src.utils import dataset_utils keras/src/utils/dataset_utils.py:9: in <module> from keras.src import tree keras/src/tree/__init__.py:1: in <module> from keras.src.tree.tree_api import assert_same_paths keras/src/tree/tree_api.py:8: in <module> from keras.src.tree import optree_impl as tree_impl keras/src/tree/optree_impl.py:4: in <module> from keras.src.backend.config import backend E ImportError: cannot import name 'backend' from partially initialized module 'keras.src.backend.config' (most likely due to a circular import) (/workspaces/keras/keras/src/backend/config.py)

So decided to revert it. 2 lines below backend is also set in the same way _BACKEND = _backend - in a way it is still consistent?

keras/src/backend/config.py

_valu

8e1c008

google-ml-butler bot added the size:M label May 5, 2025

google-ml-butler bot assigned gbaned May 5, 2025

divyashreepathihalli marked this pull request as draft May 5, 2025 23:05

divyashreepathihalli added 5 commits May 7, 2025 21:46

update variables

7159709

add nnx.jit

e378cfb

revert changes to JaxLayer

8701fc7

Merge branch 'master' into nnx

b599727

format fix

4f7b3b8

fchollet reviewed May 10, 2025

View reviewed changes

divyashreepathihalli added 7 commits May 13, 2025 20:50

make variables subclass nnx.Variable

0234e27

more tweaks

b87c4f9

update init

91b9a73

refactor jax Variable class

aee1789

code reformat

141487f

more cleanup

5ccc31e

Merge branch 'keras-team:master' into nnx

4e35416

divyashreepathihalli force-pushed the nnx branch from a1ac51b to 4e35416 Compare May 28, 2025 19:55

divyashreepathihalli added 2 commits May 28, 2025 19:57

update flax version

dd9c77d

update flax version

48983c6

divyashreepathihalli added the kokoro:force-run label May 29, 2025

kokoro-team removed the kokoro:force-run label May 29, 2025

divyashreepathihalli added 7 commits May 29, 2025 21:33

fix jax error

b22f9ef

update Variables implementation

4dbffa6

fix import

627e581

add a test

f58ef60

needs updates in operation

c2b73b7

remove __new__ from JaxLayer

a662f5e

update base optimizers

396f973

mattdangerw approved these changes Jul 24, 2025

View reviewed changes

google-ml-butler bot added kokoro:force-run ready to pull Ready to be merged into the codebase labels Jul 24, 2025

kokoro-team removed the kokoro:force-run label Jul 24, 2025

hertschuh approved these changes Jul 24, 2025

View reviewed changes

google-ml-butler bot added kokoro:force-run and removed ready to pull Ready to be merged into the codebase labels Jul 24, 2025

kokoro-team removed the kokoro:force-run label Jul 24, 2025

ypou

3635148

divyashreepathihalli force-pushed the nnx branch from b53714b to 3635148 Compare July 24, 2025 22:30

divyashreepathihalli added 2 commits July 24, 2025 23:06

resolve circukar import error

eccd843

update if condition

67f90cc

divyashreepathihalli added the kokoro:force-run label Jul 24, 2025

kokoro-team removed the kokoro:force-run label Jul 24, 2025

divyashreepathihalli added 2 commits July 25, 2025 00:03

resolve circular import error

90f3823

revert change to variables.py

a674c03

divyashreepathihalli added the kokoro:force-run label Jul 25, 2025

kokoro-team removed the kokoro:force-run label Jul 25, 2025

Merge branch 'keras-team:master' into nnx

389b19b

divyashreepathihalli added the kokoro:force-run label Jul 25, 2025

kokoro-team removed the kokoro:force-run label Jul 25, 2025

mattdangerw approved these changes Jul 25, 2025

View reviewed changes

google-ml-butler bot added kokoro:force-run ready to pull Ready to be merged into the codebase labels Jul 25, 2025

delete state comment

0426eee

google-ml-butler bot removed the ready to pull Ready to be merged into the codebase label Jul 25, 2025

kokoro-team removed the kokoro:force-run label Jul 25, 2025

divyashreepathihalli merged commit c9383e2 into keras-team:master Jul 25, 2025
8 of 11 checks passed

google-ml-butler bot removed the awaiting review label Jul 25, 2025

hertschuh mentioned this pull request Aug 8, 2025

WIP: nnx support demo #20546

Closed

Provide feedback

Saved searches

Use saved searches to filter your results more quickly

Uh oh!

Keras <> NNX integration #21252

Keras <> NNX integration #21252

Uh oh!

divyashreepathihalli commented May 5, 2025 •

edited

Loading

Uh oh!

codecov-commenter commented May 5, 2025 •

edited

Loading

Uh oh!

fchollet May 10, 2025

Uh oh!

divyashreepathihalli May 13, 2025

Uh oh!

divyashreepathihalli Jun 6, 2025

Uh oh!

mattdangerw left a comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mattdangerw Jul 24, 2025

Uh oh!

divyashreepathihalli Jul 24, 2025

Uh oh!

hertschuh Jul 24, 2025

Uh oh!

divyashreepathihalli Jul 24, 2025

Uh oh!

mattdangerw Jul 25, 2025

Uh oh!

divyashreepathihalli Jul 25, 2025

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Keras <> NNX integration #21252

Keras <> NNX integration #21252

Uh oh!

Conversation

divyashreepathihalli commented May 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Uh oh!

codecov-commenter commented May 5, 2025 • edited Loading Uh oh! There was an error while loading. Please reload this page.

Uh oh!

Codecov Report

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

mattdangerw left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

divyashreepathihalli commented May 5, 2025 •

edited

Loading

codecov-commenter commented May 5, 2025 •

edited

Loading