Keras <> NNX integration

.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,11 +20,17 @@ 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 }}
       KERAS_HOME: .github/workflows/config/${{ matrix.backend }}
+      KERAS_BACKEND: ${{ matrix.backend }}
     steps:
       - uses: actions/checkout@v4
       - name: Check for changes in keras/src/applications
@@ -48,15 +57,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 +77,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 +103,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 +114,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_BACKEND,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
 }

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

@@ -0,0 +1,7 @@
+{
+    "floatx": "float32",
+    "epsilon": 1e-07,
+    "backend": "jax",
+    "image_data_format": "channels_last",
+    "nnx_enabled": true
+} 

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 = {
@@ -55,16 +56,25 @@ def manage_venv_installs(whl_path):
         "pip install " + backend_extra_url + backend_pkg,
         "pip install -r requirements-common.txt",
         "pip install pytest",
-        # Ensure other backends are uninstalled
-        "pip uninstall -y "
-        + BACKEND_REQ[other_backends[0]][0]
-        + " "
-        + BACKEND_REQ[other_backends[1]][0]
-        + " "
-        + BACKEND_REQ[other_backends[2]][0],
-        # 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")
+
+    install_setup.extend(
+        [
+            # Ensure other backends are uninstalled
+            "pip uninstall -y "
+            + BACKEND_REQ[other_backends[0]][0]
+            + " "
+            + BACKEND_REQ[other_backends[1]][0]
+            + " "
+            + BACKEND_REQ[other_backends[2]][0],
+            # Install `.whl` package
+            "pip install " + whl_path,
+        ]
+    )
     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,9 @@ 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 becomes a no-op for this variable.
             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()
@@ -261,6 +290,30 @@ def keras_home():
 
 # Attempt to read Keras config file.
 _config_path = os.path.expanduser(os.path.join(_KERAS_DIR, "keras.json"))
+
+# Save config file, if possible.
+if not os.path.exists(_KERAS_DIR):
+    try:
+        os.makedirs(_KERAS_DIR)
+    except OSError:
+        # Except permission denied.
+        pass
+
+if not os.path.exists(_config_path):
+    _config_to_save = {
+        "floatx": _FLOATX,
+        "epsilon": _EPSILON,
+        "backend": _BACKEND,
+        "image_data_format": _IMAGE_DATA_FORMAT,
+        "nnx_enabled": is_nnx_enabled(),
+    }
+    try:
+        with open(_config_path, "w") as f:
+            f.write(json.dumps(_config_to_save, indent=4))
+    except IOError:
+        # Except permission denied.
+        pass
+
 if os.path.exists(_config_path):
     try:
         with open(_config_path) as f:
@@ -274,36 +327,17 @@ 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)
+    if isinstance(_nnx_enabled_config, bool):
+        _NNX_ENABLED = _nnx_enabled_config
+    # else: ignore non-bool values for nnx_enabled
 
+    # Apply basic configs that don't cause circular import
     set_floatx(_floatx)
     set_epsilon(_epsilon)
     set_image_data_format(_image_data_format)
     _BACKEND = _backend
 
-# Save config file, if possible.
-if not os.path.exists(_KERAS_DIR):
-    try:
-        os.makedirs(_KERAS_DIR)
-    except OSError:
-        # Except permission denied and potential race conditions
-        # in multi-threaded environments.
-        pass
-
-if not os.path.exists(_config_path):
-    _config = {
-        "floatx": floatx(),
-        "epsilon": epsilon(),
-        "backend": _BACKEND,
-        "image_data_format": image_data_format(),
-    }
-    try:
-        with open(_config_path, "w") as f:
-            f.write(json.dumps(_config, indent=4))
-    except IOError:
-        # Except permission denied.
-        pass
-
-# Set backend based on KERAS_BACKEND flag, if applicable.
 if "KERAS_BACKEND" in os.environ:
     _backend = os.environ["KERAS_BACKEND"]
     if _backend:
@@ -313,6 +347,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 +438,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:
+        _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