Keras <> NNX integration

.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]
     name: Run tests
     runs-on: ubuntu-latest
     env:

.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
+}

guides/distributed_training_with_jax.py

@@ -53,6 +53,9 @@
 from jax.sharding import Mesh
 from jax.sharding import NamedSharding
 from jax.sharding import PartitionSpec as P
+from keras.src.backend.config import is_nnx_enabled
+from keras.src.utils.jax_utils import jit
+from flax import nnx
 
 
 def get_model():
@@ -186,7 +189,7 @@ def compute_loss(trainable_variables, non_trainable_variables, x, y):
 
 
 # Training step, Keras provides a pure functional optimizer.stateless_apply
-@jax.jit
+@jit
 def train_step(train_state, x, y):
     (
         trainable_variables,

integration_tests/import_test.py

@@ -11,7 +11,9 @@
         "torch",
         "--extra-index-url https://download.pytorch.org/whl/cpu ",
     ),
-    "jax": ("jax[cpu]", ""),
+    # please update the jax version here if jax version is updated in
+    # requirements file
+    "jax": ("jax[cpu]==0.5.0", ""),
     "openvino": ("openvino", ""),
 }
 

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/__init__.py

@@ -39,7 +39,7 @@
     from keras.src.backend.tensorflow.core import Variable as BackendVariable
 elif backend() == "jax":
     from keras.src.backend.jax import *  # noqa: F403
-    from keras.src.backend.jax.core import Variable as BackendVariable
+    from keras.src.backend.jax import Variable as BackendVariable
 elif backend() == "torch":
     from keras.src.backend.torch import *  # noqa: F403
     from keras.src.backend.torch.core import Variable as BackendVariable

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,33 @@ 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 the NNX backend, you must install `flax`."
+                "Try: `pip 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 +291,7 @@ def keras_home():
 
 # Attempt to read Keras config file.
 _config_path = os.path.expanduser(os.path.join(_KERAS_DIR, "keras.json"))
+
 if os.path.exists(_config_path):
     try:
         with open(_config_path) as f:
@@ -274,36 +305,18 @@ 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 not isinstance(_nnx_enabled_config, bool):
+        _NNX_ENABLED = str(_nnx_enabled_config).lower() == "true"
+    else:
+        _NNX_ENABLED = _nnx_enabled_config
 
+    # 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 +326,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 +417,35 @@ def max_steps_per_epoch():
             `None`, no limit is applied.
     """
     return _MAX_STEPS_PER_EPOCH
+
+
+if not os.path.exists(_KERAS_DIR):
+    try:
+        os.makedirs(_KERAS_DIR)
+    except OSError:
+        # Except permission denied and potential race conditions
+        pass
+
+if not os.path.exists(_config_path):
+    _config_to_save = {
+        "floatx": floatx(),
+        "epsilon": epsilon(),
+        "backend": _BACKEND,  # Use the final _BACKEND value
+        "image_data_format": image_data_format(),
+        "nnx_enabled": _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 "KERAS_NNX_ENABLED" in os.environ:
+    env_val = os.environ["KERAS_NNX_ENABLED"].lower()
+    if env_val == "true":
+        _NNX_ENABLED = True
+    elif env_val == "false":
+        _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
@@ -10,7 +11,11 @@
 from keras.src.backend.jax.core import IS_THREAD_SAFE
 from keras.src.backend.jax.core import SUPPORTS_RAGGED_TENSORS
 from keras.src.backend.jax.core import SUPPORTS_SPARSE_TENSORS
-from keras.src.backend.jax.core import Variable
+
+if is_nnx_enabled():
+    from keras.src.backend.jax.core import NnxVariable as Variable
+else:
+    from keras.src.backend.jax.core import JaxVariable as Variable
 from keras.src.backend.jax.core import cast
 from keras.src.backend.jax.core import compute_output_spec
 from keras.src.backend.jax.core import cond