simplify init

Author: divyashreepathihalli

keras/src/backend/jax/__init__.py

@@ -11,11 +11,7 @@
 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
-
-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 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

keras/src/backend/jax/core.py

@@ -58,7 +58,7 @@ def __jax_array__(self):
         return self.value
 
 
-_JAX_VARIABLE_TYPE = JaxVariable
+Variable = JaxVariable
 if config.is_nnx_enabled():
     from flax import nnx
 
@@ -231,7 +231,7 @@ def value(self):
                 )
             return self._maybe_autocast(current_value)
 
-    _JAX_VARIABLE_TYPE = NnxVariable
+    Variable = NnxVariable
 
 
 def convert_to_tensor(x, dtype=None, sparse=None, ragged=None):
@@ -247,7 +247,7 @@ def convert_to_tensor(x, dtype=None, sparse=None, ragged=None):
         # an existing distributed jax array will raise error.
         return x
 
-    if isinstance(x, _JAX_VARIABLE_TYPE):
+    if isinstance(x, Variable):
         if dtype is not None and x.dtype != dtype:
             return x.value.astype(dtype)
         return x.value
@@ -531,7 +531,7 @@ def fori_loop(lower, upper, body_fun, init_val):
 
 
 def stop_gradient(variable):
-    if isinstance(variable, _JAX_VARIABLE_TYPE):
+    if isinstance(variable, Variable):
         variable = variable.value
     return jax.lax.stop_gradient(variable)