fix import error

Author: divyashreepathihalli

keras/src/backend/jax/core.py

@@ -1,9 +1,9 @@
+import flax
 import jax
 import jax.experimental.sparse as jax_sparse
 import jax.numpy as jnp
 import ml_dtypes
 import numpy as np
-from flax import nnx
 
 from keras.src import tree
 from keras.src.backend.common import KerasVariable
@@ -57,7 +57,7 @@ def __jax_array__(self):
         return self.value
 
 
-class Variable(JaxVariable, nnx.Variable):
+class Variable(JaxVariable, flax.nnx.Variable):
     def __init__(
         self,
         initializer,
@@ -123,7 +123,7 @@ def _complete_nnx_init(self):
             current_nnx_mutable = self.trainable  # A sensible default link
 
         # initialize the nnx.Variable
-        nnx.Variable.__init__(
+        flax.nnx.Variable.__init__(
             self,
             value=self._value,
             mutable=current_nnx_mutable,
@@ -173,8 +173,8 @@ def value(self, new_value):
     # Overriding NNX methods that modify `raw_value` or `_var_metadata` directly
     # to ensure Keras's `_value` and other Keras states are in sync.
 
-    def copy_from(self, other: nnx.Variable):  # type: ignore
-        if not isinstance(other, nnx.Variable):  # Basic check from nnx
+    def copy_from(self, other: flax.nnx.Variable):  # type: ignore
+        if not isinstance(other, flax.nnx.Variable):  # Basic check from nnx
             raise TypeError(
                 f"Expected nnx.Variable, got {type(other).__name__}"
             )
@@ -184,12 +184,12 @@ def copy_from(self, other: nnx.Variable):  # type: ignore
         # Let nnx.Variable handle its part (updates self.raw_value and
         # self._var_metadata)
         # Need to call nnx.Variable.copy_from specifically.
-        nnx.Variable.copy_from(self, other)
+        flax.nnx.Variable.copy_from(self, other)
 
         # Now, self.raw_value is updated. Sync Keras's self._value.
         # Extract the JAX array if raw_value is a nnx.mutable_array
         keras_value_to_assign = self.raw_value
-        if nnx.utils.is_mutable_array(keras_value_to_assign):
+        if flax.nnx.utils.is_mutable_array(keras_value_to_assign):
             keras_value_to_assign = keras_value_to_assign.__array__()
 
         self.assign(keras_value_to_assign)
@@ -205,11 +205,11 @@ def copy_from(self, other: nnx.Variable):  # type: ignore
     def update_from_state(self, variable_state: nnx.graph.VariableState):
         # Let nnx.Variable handle its part (updates self.raw_value and
         # self._var_metadata)
-        nnx.Variable.update_from_state(self, variable_state)
+        flax.nnx.Variable.update_from_state(self, variable_state)
 
         # Sync Keras's self._value
         keras_value_to_assign = self.raw_value
-        if nnx.utils.is_mutable_array(keras_value_to_assign):
+        if flax.nnx.utils.is_mutable_array(keras_value_to_assign):
             keras_value_to_assign = keras_value_to_assign.__array__()
 
         self.assign(keras_value_to_assign)
@@ -242,7 +242,7 @@ def __getstate__(self):
             "_nnx_metadata_arg": self._nnx_metadata_arg,
             "_nnx_init_pending": self._nnx_init_pending,
         }
-        nnx_state = nnx.Variable.__getstate__(self)
+        nnx_state = flax.nnx.Variable.__getstate__(self)
         return {"keras_state": keras_state, "nnx_state": nnx_state}
 
     def __setstate__(self, state):
@@ -254,7 +254,7 @@ def __setstate__(self, state):
             object.__setattr__(self, k, v)
 
         # Restore NNX attributes using its __setstate__
-        nnx.Variable.__setstate__(self, nnx_state)
+        flax.nnx.Variable.__setstate__(self, nnx_state)
 
         if (
             self._initializer is not None and self._value is None
@@ -276,7 +276,7 @@ def __setstate__(self, state):
                 if current_nnx_mutable is None:
                     current_nnx_mutable = self.trainable
 
-                if current_nnx_mutable and nnx.utils.is_mutable_array(
+                if current_nnx_mutable and flax.nnx.utils.is_mutable_array(
                     self.raw_value
                 ):
                     self.raw_value[...] = self._value