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 15.26316% with 161 lines in your changes missing coverage. Please review.

Project coverage is 77.57%. Comparing base (81821e0) to head (c05166e).
Report is 18 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/src/backend/jax/core.py	4.46%	105 Missing and 2 partials ⚠️
keras/src/backend/config.py	30.55%	20 Missing and 5 partials ⚠️
keras/src/layers/layer.py	30.76%	8 Missing and 1 partial ⚠️
keras/src/utils/jax_utils.py	22.22%	7 Missing ⚠️
keras/src/backend/jax/trainer.py	0.00%	5 Missing ⚠️
keras/src/backend/jax/layer.py	0.00%	4 Missing ⚠️
keras/src/backend/jax/__init__.py	50.00%	1 Missing and 1 partial ⚠️
keras/api/_tf_keras/keras/config/__init__.py	0.00%	1 Missing ⚠️
keras/src/backend/__init__.py	0.00%	1 Missing ⚠️

❗ There is a different number of reports uploaded between BASE (81821e0) and HEAD (c05166e). Click for more details.

HEAD has 2 uploads less than BASE

Flag BASE (81821e0) HEAD (c05166e)

keras 5 4

keras-jax 1 0

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #21252      +/-   ##
==========================================
- Coverage   82.65%   77.57%   -5.09%     
==========================================
  Files         565      565              
  Lines       54802    55062     +260     
  Branches     8508     8552      +44     
==========================================
- Hits        45297    42714    -2583     
- Misses       7413    10337    +2924     
+ Partials     2092     2011      -81

Flag	Coverage Δ
keras	`77.41% <15.26%> (-5.05%)`	⬇️
keras-jax	`?`
keras-numpy	`58.59% <15.26%> (-0.16%)`	⬇️
keras-openvino	`33.48% <14.21%> (+0.33%)`	⬆️
keras-tensorflow	`63.81% <15.26%> (-0.19%)`	⬇️
keras-torch	`63.45% <15.26%> (-0.20%)`	⬇️

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

        import jax.numpy as jnp

        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"]

…rting

keras/src/ops/operation.py

keras/src/backend/jax/trainer.py

keras/src/backend/jax/core.py

keras/src/backend/jax/trainer.py

requirements-jax-cuda.txt

requirements.txt

mattdangerw

This needs a testing plan

How do we make sure things work without flax installed on CI?
How do we make sure things work when NNX is enabled?
When NNX is not enabled?

We need to make sure we have an automated testing path for these different options we are writing logic for or they will silently break at some point.

keras/src/backend/config.py

keras/src/backend/jax/core.py

divyashreepathihalli · 2025-06-06T19:14:35Z

Added a github workflow action for nnx backend. Note this will FAIL - because this needs a new release of flax to work.

hertschuh

Thanks for the PR!

requirements-common.txt

requirements-jax-cuda.txt

requirements.txt

hertschuh · 2025-06-06T20:51:26Z

.github/workflows/actions.yml

+      - name: Install Flax for NNX backend
+        if: matrix.backend == 'nnx'
+        run: |
+          pip install flax --progress-bar off --upgrade


Not needed, it's already part of the requirements file.

hertschuh · 2025-06-06T20:52:35Z

.github/workflows/actions.yml

@@ -16,7 +16,7 @@ jobs:
      fail-fast: false
      matrix:
        python-version: ['3.10']
-        backend: [tensorflow, jax, torch, numpy, openvino]
+        backend: [tensorflow, jax, torch, numpy, openvino, nnx]


I think you need to do something else in addition (or instead of this?), right now, this does not turn on NNX.

hertschuh · 2025-06-06T21:11:47Z

keras/src/backend/jax/core.py

+                elif "mutable" not in nnx_metadata:
+                    nnx_metadata["mutable"] = actual_nnx_mutable
+
+                # Initialize nnx.Variable first.


Do we need a placeholder value? Does it allocate HBM? The issue is that we may be using twice the memory we need by initializing the variable twice.

hertschuh · 2025-06-06T21:12:09Z

keras/src/backend/jax/core.py

+                else:
+                    _placeholder_value = jnp.array(0.0, dtype=jnp.float32)
+
+                # Call nnx.Variable.__init__ directly.


What about the sharding on the placeholder_value, do we need it?

hertschuh · 2025-06-06T21:15:36Z

keras/src/backend/jax/core.py

+
+                return self._maybe_autocast(current_value)
+
+            def __hash__(self):


Per the discussion we had, we can't do that. It's not compatible with the semantics of __eq__ that we have.

I think my preferred fix would be to change NNX to not use variables as keys, but instead use id(variable) in dictionaries.

keras/src/ops/operation.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

divyashreepathihalli added 13 commits June 2, 2025 20:37

update variable value updates

b99571a

sync value properly

ed0bc00

update flag based routing between nnx and jax

460e0e2

clean up

34f27e9

fix circular import error

427ff82

fix is nnx call enabled flag

c4ee191

attemptto fix circular import error

44414dc

try again

0953d99

fix import error

6d54a7e

reformat# Please enter the commit message for your changes. Lines sta…

64adbaf

…rting

This has to fix it

6454800

api gen

001f112

remove enable diisable configs -that does not work

5f26958

divyashreepathihalli marked this pull request as ready for review June 4, 2025 05:56

mattdangerw reviewed Jun 5, 2025

View reviewed changes

keras/src/ops/operation.py Outdated Show resolved Hide resolved

keras/src/backend/jax/trainer.py Outdated Show resolved Hide resolved

keras/src/backend/jax/core.py Outdated Show resolved Hide resolved

keras/src/backend/jax/trainer.py Outdated Show resolved Hide resolved

mattdangerw reviewed Jun 5, 2025

View reviewed changes

requirements-jax-cuda.txt Show resolved Hide resolved

mattdangerw reviewed Jun 5, 2025

View reviewed changes

requirements.txt Show resolved Hide resolved

mattdangerw reviewed Jun 5, 2025

View reviewed changes

keras/src/backend/config.py Outdated Show resolved Hide resolved

keras/src/backend/config.py Outdated Show resolved Hide resolved

keras/src/backend/config.py Outdated Show resolved Hide resolved

keras/src/backend/jax/core.py Outdated Show resolved Hide resolved

divyashreepathihalli added 4 commits June 6, 2025 18:01

adrress some comments

782c653

update conditional imports

561f70a

fix tests

e7caa03

add github workflow for nnx

8e3f460

fix test

d70d51c

hertschuh reviewed Jun 6, 2025

View reviewed changes

divyashreepathihalli added 5 commits June 6, 2025 22:42

address comments

38dbd4b

fix test

1c60c5e

address comments

74835fd

fix test

6f11c0c

fix test -_-

c05166e


		return self._maybe_autocast(current_value)

		def __hash__(self):

Keras <> NNX integration #21252

Are you sure you want to change the base?

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!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

mattdangerw left a comment

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

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

Uh oh!

Uh oh!

hertschuh left a comment

Choose a reason for hiding this comment

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!

Uh oh!

Uh oh!

Uh oh!

divyashreepathihalli commented May 5, 2025 •

edited

Loading

codecov-commenter commented May 5, 2025 •

edited

Loading

divyashreepathihalli commented Jun 6, 2025 •

edited

Loading