Migration to Orbax V1

Author: amitsrivastava78

keras/api/_tf_keras/keras/distribution/__init__.py

@@ -17,7 +17,6 @@
 from keras.src.distribution.distribution_lib import distribution as distribution
 from keras.src.distribution.distribution_lib import initialize as initialize
 from keras.src.distribution.distribution_lib import list_devices as list_devices
-from keras.src.distribution.distribution_lib import process_id as process_id
 from keras.src.distribution.distribution_lib import (
     set_distribution as set_distribution,
 )

keras/api/distribution/__init__.py

@@ -17,7 +17,6 @@
 from keras.src.distribution.distribution_lib import distribution as distribution
 from keras.src.distribution.distribution_lib import initialize as initialize
 from keras.src.distribution.distribution_lib import list_devices as list_devices
-from keras.src.distribution.distribution_lib import process_id as process_id
 from keras.src.distribution.distribution_lib import (
     set_distribution as set_distribution,
 )

keras/src/backend/__init__.py

@@ -36,9 +36,11 @@
 # Import backend functions.
 if backend() == "tensorflow":
     from keras.src.backend.tensorflow import *  # noqa: F403
+    from keras.src.backend.tensorflow import distribution_lib
     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 import distribution_lib
     from keras.src.backend.jax.core import Variable as BackendVariable
 elif backend() == "torch":
     from keras.src.backend.torch import *  # noqa: F403

keras/src/backend/jax/__init__.py

@@ -14,7 +14,6 @@
 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
-from keras.src.backend.jax.core import convert_checkpoint_value
 from keras.src.backend.jax.core import convert_to_numpy
 from keras.src.backend.jax.core import convert_to_tensor
 from keras.src.backend.jax.core import device_scope
@@ -25,7 +24,6 @@
 from keras.src.backend.jax.core import shape
 from keras.src.backend.jax.core import stop_gradient
 from keras.src.backend.jax.core import vectorized_map
-from keras.src.backend.jax.distribution_lib import process_id
 from keras.src.backend.jax.rnn import cudnn_ok
 from keras.src.backend.jax.rnn import gru
 from keras.src.backend.jax.rnn import lstm

keras/src/backend/jax/core.py

@@ -572,35 +572,3 @@ def device_scope(device_name):
     else:
         jax_device = device_name
     return jax.default_device(jax_device)
-
-
-def convert_checkpoint_value(value, dtype, shape):
-    """Convert a value for checkpoint restoration, preserving JAX arrays for
-    sharding.
-
-    This function handles the special case of checkpoint restoration where JAX
-    arrays should be preserved for sharding support, while other values are
-    converted to JAX arrays with the specified dtype and shape.
-
-    Args:
-        value: The value to convert (can be JAX array, numpy array, or other
-            types)
-        dtype: The target dtype
-        shape: The target shape
-
-    Returns:
-        A JAX array with the specified dtype and shape, or the original JAX
-        array if it was already a JAX array.
-    """
-    # For JAX backend, preserve JAX arrays for sharding support
-    if hasattr(value, "__array_namespace__") or str(type(value)).startswith(
-        "<class 'jax"
-    ):
-        # value is already a JAX array, return as-is to preserve sharding
-        return value
-    elif isinstance(value, np.ndarray):
-        # Convert numpy array to JAX array
-        return jnp.array(value).astype(dtype).reshape(shape)
-    else:
-        # Convert other types to JAX array
-        return jnp.array(value, dtype=dtype).reshape(shape)

keras/src/backend/jax/distribution_lib.py

@@ -193,11 +193,6 @@ def num_processes():
     return jax.process_count()
 
 
-def process_id():
-    """Return the current process ID for the distribution setting."""
-    return jax.process_index()
-
-
 def _to_backend_device(device_name):
     if isinstance(device_name, jax.Device):
         return device_name

keras/src/backend/numpy/__init__.py

@@ -20,8 +20,10 @@
 from keras.src.backend.numpy.core import random_seed_dtype
 from keras.src.backend.numpy.core import shape
 from keras.src.backend.numpy.core import vectorized_map
-from keras.src.backend.numpy.distribution_lib import process_id
 from keras.src.backend.numpy.rnn import cudnn_ok
 from keras.src.backend.numpy.rnn import gru
 from keras.src.backend.numpy.rnn import lstm
 from keras.src.backend.numpy.rnn import rnn
+
+# Numpy backend does not support distribution
+distribution_lib = None

keras/src/backend/numpy/distribution_lib.py

@@ -1,6 +1,5 @@
 from keras.src.backend.common.name_scope import name_scope
 from keras.src.backend.openvino import core
-from keras.src.backend.openvino import distribution_lib
 from keras.src.backend.openvino import image
 from keras.src.backend.openvino import linalg
 from keras.src.backend.openvino import math
@@ -20,7 +19,6 @@
 from keras.src.backend.openvino.core import random_seed_dtype
 from keras.src.backend.openvino.core import shape
 from keras.src.backend.openvino.core import vectorized_map
-from keras.src.backend.openvino.distribution_lib import process_id
 from keras.src.backend.openvino.rnn import cudnn_ok
 from keras.src.backend.openvino.rnn import gru
 from keras.src.backend.openvino.rnn import lstm