Keras <> NNX integration (#21252)

* _valu

* update variables

* add nnx.jit

* revert changes to JaxLayer

* format fix

* make variables subclass nnx.Variable

* more tweaks

* update init

* refactor jax Variable class

* code reformat

* more cleanup

* update flax version

* update flax version

* fix jax error

* update Variables implementation

* fix import

* add a test

* needs updates in operation

* remove __new__ from JaxLayer

* update base optimizers

* code reformat+ model saving tests

* add __hash__

* update variable value updates

* sync value properly

* update flag based routing between nnx and jax

* clean up

* fix circular import error

* fix is nnx call enabled flag

* attemptto fix circular import error

* try again

* fix import error

* reformat# Please enter the commit message for your changes. Lines starting

* This has to fix it

* api gen

* remove enable diisable configs -that does not work

* adrress some comments

* update conditional imports

* fix tests

* add github workflow for nnx

* fix test

* address comments

* fix test

* address comments

* fix test

* fix test -_-

* put the set attr in operation

* fix jax error

* fix trace error

* remove installation

* import fixes

* update jax version

* ugh the jax version issue

* update jax version

* update installations

* update jax utils

* another requirents file fix

* fix test

* add back flax to req common

* address review comments

* fix tests

* fix tests address more comments

* fix tests

* fix tests

* fix jax tests

* revert guide

* fix code format

* fix tests and jit

* fix import test

* try to fix memory error

* revert memory fix

* fix test

* fix test

* fix test

* revert version back

* Update functional.py

* Update core.py

* remove nnx workflow

* code reformat

* address gemini comments

* address latest comments

* remove hash function

* update tests name

* address latest comments

* address review comments

* fix actions

* fix actions

* skipt other tests in nnx backend

* revert changes to basic flow

* point installation to official JAX code

* fix actions

* revert basic flow test

* move logic out of functional and into function class

* revert functional.py

* simplify init

* FIX MODEL BUILD ERROR

* revert changes to basic_full_flow.py

* revert basic_full_flow.py

* address review comments

* add layer.py#	modified:   keras/src/layers/layer.py

* revert variables change

* nit

* update init

* ypou

* resolve circukar import error

* update if condition

* resolve circular import error

* revert change to variables.py

* delete state comment

Author: divyashreepathihalli

.github/workflows/actions.yml

@@ -1,5 +1,8 @@
 name: Tests
 
+# TODO: Consider enabling all tests (pytest, applications, etc.) with NNX in the future
+# Currently only basic flow tests run with NNX enabled
+
 on:
   push:
     branches: [ master ]
@@ -17,7 +20,12 @@ jobs:
       matrix:
         python-version: ['3.10']
         backend: [tensorflow, jax, torch, numpy, openvino]
-    name: Run tests
+        nnx_enabled: [false]
+        include:
+          - python-version: '3.10'
+            backend: jax
+            nnx_enabled: true
+    name: ${{ matrix.backend == 'jax' && format('Run tests ({0}, {1}, nnx_enabled = {2})', matrix.python-version, matrix.backend, matrix.nnx_enabled) || format('Run tests ({0}, {1})', matrix.python-version, matrix.backend) }}
     runs-on: ubuntu-latest
     env:
       PYTHON: ${{ matrix.python-version }}
@@ -48,15 +56,18 @@ jobs:
       - name: Install dependencies
         run: |
           pip install -r requirements.txt --progress-bar off --upgrade
+          if [ "${{ matrix.nnx_enabled }}" == "true" ]; then
+            pip install --upgrade git+https://github.com/google/flax.git
+          fi
           pip uninstall -y keras keras-nightly
           pip install -e "." --progress-bar off --upgrade
       - name: Test applications with pytest
-        if: ${{ steps.filter.outputs.applications == 'true' }}
+        if: ${{ steps.filter.outputs.applications == 'true' && matrix.nnx_enabled == false }}
         run: |
           pytest keras/src/applications --cov=keras/src/applications --cov-config=pyproject.toml
           coverage xml --include='keras/src/applications/*' -o apps-coverage.xml
       - name: Codecov keras.applications
-        if: ${{ steps.filter.outputs.applications == 'true' }}
+        if: ${{ steps.filter.outputs.applications == 'true' && matrix.nnx_enabled == false }}
         uses: codecov/codecov-action@v5
         with:
           env_vars: PYTHON,KERAS_HOME
@@ -65,14 +76,21 @@ jobs:
           token: ${{ secrets.CODECOV_TOKEN }}
           fail_ci_if_error: false
       - name: Test integrations
-        if: ${{ matrix.backend != 'numpy'}}
+        if: ${{ matrix.backend != 'numpy' && matrix.nnx_enabled == false }}
         run: |
           python integration_tests/import_test.py
           python integration_tests/numerical_test.py
       - name: Test JAX-specific integrations
-        if: ${{ matrix.backend == 'jax'}}
+        if: ${{ matrix.backend == 'jax' && matrix.nnx_enabled == false }}
         run: |
           python integration_tests/jax_custom_fit_test.py
+      - name: Test basic flow with NNX
+        if: ${{ matrix.nnx_enabled == true }}
+        env:
+          KERAS_NNX_ENABLED: true
+        run: |
+          python integration_tests/import_test.py
+          python integration_tests/basic_full_flow.py
       - name: Test TF-specific integrations
         if: ${{ matrix.backend == 'tensorflow'}}
         run: |
@@ -84,6 +102,7 @@ jobs:
           pytest integration_tests/torch_workflow_test.py
           python integration_tests/torch_custom_fit_test.py
       - name: Test with pytest
+        if: ${{ matrix.nnx_enabled == false }}
         run: |
           if [ "${{ matrix.backend }}" == "openvino" ]; then
             IGNORE_FILE="keras/src/backend/openvino/excluded_tests.txt"
@@ -94,10 +113,11 @@ jobs:
           pytest keras --ignore keras/src/applications --cov=keras --cov-config=pyproject.toml $IGNORE_ARGS
           coverage xml --omit='keras/src/applications/*,keras/api' -o core-coverage.xml
       - name: Codecov keras
+        if: ${{ matrix.nnx_enabled == false }}
         uses: codecov/codecov-action@v5
         with:
-          env_vars: PYTHON,KERAS_HOME
-          flags: keras,keras-${{ matrix.backend }}
+          env_vars: PYTHON,KERAS_HOME,KERAS_NNX_ENABLED
+          flags: keras,keras-${{ matrix.backend }}${{ matrix.nnx_enabled == 'true' && '-nnx' || '' }}
           files: core-coverage.xml
           token: ${{ secrets.CODECOV_TOKEN }}
           fail_ci_if_error: false

.github/workflows/config/jax/keras.json

@@ -2,5 +2,6 @@
     "floatx": "float32",
     "epsilon": 1e-07,
     "backend": "jax",
-    "image_data_format": "channels_last"
+    "image_data_format": "channels_last",
+    "nnx_enabled": false
 }

integration_tests/import_test.py

@@ -3,6 +3,7 @@
 import subprocess
 
 from keras.src import backend
+from keras.src.backend import config
 
 # For torch, use index url to avoid installing nvidia drivers for the test.
 BACKEND_REQ = {
@@ -65,6 +66,9 @@ def manage_venv_installs(whl_path):
         # Install `.whl` package
         "pip install " + whl_path,
     ]
+    # Install flax for JAX when NNX is enabled
+    if backend.backend() == "jax" and config.is_nnx_enabled():
+        install_setup.append("pip install flax>=0.10.1")
     run_commands_venv(install_setup)
 
 

keras/api/_tf_keras/keras/config/__init__.py

@@ -17,6 +17,7 @@
 from keras.src.backend.config import (
     is_flash_attention_enabled as is_flash_attention_enabled,
 )
+from keras.src.backend.config import is_nnx_enabled as is_nnx_enabled
 from keras.src.backend.config import max_epochs as max_epochs
 from keras.src.backend.config import max_steps_per_epoch as max_steps_per_epoch
 from keras.src.backend.config import set_epsilon as set_epsilon

keras/api/config/__init__.py

@@ -17,6 +17,7 @@
 from keras.src.backend.config import (
     is_flash_attention_enabled as is_flash_attention_enabled,
 )
+from keras.src.backend.config import is_nnx_enabled as is_nnx_enabled
 from keras.src.backend.config import max_epochs as max_epochs
 from keras.src.backend.config import max_steps_per_epoch as max_steps_per_epoch
 from keras.src.backend.config import set_epsilon as set_epsilon

keras/src/backend/common/variables.py

@@ -211,6 +211,11 @@ def __init__(
 
     def _deferred_initialize(self):
         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
+            # returns early and does not raise an error.
+            if config.is_nnx_enabled():
+                return
             raise ValueError(f"Variable {self.path} is already initialized.")
 
         if in_stateless_scope():

keras/src/backend/config.py

@@ -15,6 +15,9 @@
 # Default backend: TensorFlow.
 _BACKEND = "tensorflow"
 
+# Whether NNX is enabled.
+_NNX_ENABLED = False
+
 # Cap run duration for debugging.
 _MAX_EPOCHS = None
 _MAX_STEPS_PER_EPOCH = None
@@ -230,6 +233,32 @@ def is_flash_attention_enabled():
     return global_state.get_global_attribute("flash_attention", default=None)
 
 
+@keras_export("keras.config.is_nnx_enabled")
+def is_nnx_enabled():
+    """Checks whether NNX specific features are enabled for the JAX backend.
+
+    Returns:
+        bool: `True` if NNX backend features are enabled, `False` otherwise.
+        Defaults to `False`.
+    """
+    return _NNX_ENABLED
+
+
+def set_nnx_enabled(value):
+    global _NNX_ENABLED
+    from keras.src.backend.common import global_state
+
+    _NNX_ENABLED = bool(value)
+    if _NNX_ENABLED:
+        try:
+            from flax import nnx  # noqa F401
+        except ImportError:
+            raise ImportError(
+                "To use NNX with the JAX backend, you must install `flax`."
+            )
+    global_state.set_global_attribute("nnx_enabled", bool(value))
+
+
 def standardize_data_format(data_format):
     if data_format is None:
         return image_data_format()
@@ -274,8 +303,11 @@ def keras_home():
     _backend = _config.get("backend", _BACKEND)
     _image_data_format = _config.get("image_data_format", image_data_format())
     assert _image_data_format in {"channels_last", "channels_first"}
+    _nnx_enabled_config = _config.get("nnx_enabled", _NNX_ENABLED)
 
+    # Apply basic configs that don't cause circular import
     set_floatx(_floatx)
+    _NNX_ENABLED = _nnx_enabled_config
     set_epsilon(_epsilon)
     set_image_data_format(_image_data_format)
     _BACKEND = _backend
@@ -313,6 +345,7 @@ def keras_home():
 if "KERAS_MAX_STEPS_PER_EPOCH" in os.environ:
     _MAX_STEPS_PER_EPOCH = int(os.environ["KERAS_MAX_STEPS_PER_EPOCH"])
 
+
 if _BACKEND != "tensorflow":
     # If we are not running on the tensorflow backend, we should stop tensorflow
     # from using all available GPU memory. See
@@ -403,3 +436,13 @@ def max_steps_per_epoch():
             `None`, no limit is applied.
     """
     return _MAX_STEPS_PER_EPOCH
+
+
+if "KERAS_NNX_ENABLED" in os.environ:
+    env_val = os.environ["KERAS_NNX_ENABLED"].lower()
+    if env_val == "true" or env_val == "1":
+        _NNX_ENABLED = True
+    else:
+        _NNX_ENABLED = False
+
+set_nnx_enabled(_NNX_ENABLED)

keras/src/backend/jax/__init__.py

@@ -1,3 +1,4 @@
+from keras.src.backend.config import is_nnx_enabled
 from keras.src.backend.jax import core
 from keras.src.backend.jax import distribution_lib
 from keras.src.backend.jax import image