Migration to Orbax V1

Author: amitsrivastava78

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

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

keras/src/backend/tensorflow/__init__.py

@@ -13,7 +13,6 @@
 from keras.src.backend.tensorflow.core import cast
 from keras.src.backend.tensorflow.core import compute_output_spec
 from keras.src.backend.tensorflow.core import cond
-from keras.src.backend.tensorflow.core import convert_checkpoint_value
 from keras.src.backend.tensorflow.core import convert_to_numpy
 from keras.src.backend.tensorflow.core import convert_to_tensor
 from keras.src.backend.tensorflow.core import device_scope

keras/src/backend/tensorflow/core.py

@@ -696,23 +696,3 @@ def __exit__(self, *args, **kwargs):
 
 def device_scope(device_name):
     return tf.device(device_name)
-
-
-def convert_checkpoint_value(value, dtype, shape):
-    """Convert a value for checkpoint restoration.
-
-    For TensorFlow backend, convert to numpy arrays with specified dtype and
-    shape.
-
-    Args:
-        value: The value to convert
-        dtype: The target dtype
-        shape: The target shape
-
-    Returns:
-        A numpy array with the specified dtype and shape.
-    """
-    if isinstance(value, np.ndarray):
-        return value.astype(dtype).reshape(shape)
-    else:
-        return np.array(value, dtype=dtype).reshape(shape)

keras/src/backend/torch/__init__.py

@@ -29,7 +29,6 @@
 from keras.src.backend.torch.core import cast
 from keras.src.backend.torch.core import compute_output_spec
 from keras.src.backend.torch.core import cond
-from keras.src.backend.torch.core import convert_checkpoint_value
 from keras.src.backend.torch.core import convert_to_numpy
 from keras.src.backend.torch.core import convert_to_tensor
 from keras.src.backend.torch.core import device_scope

keras/src/backend/torch/core.py

@@ -730,22 +730,3 @@ def backward(ctx, grad_output):
         if not isinstance(grads, tuple):
             grads = (grads,)
         return (None,) + grads
-
-
-def convert_checkpoint_value(value, dtype, shape):
-    """Convert a value for checkpoint restoration.
-
-    For PyTorch backend, convert to numpy arrays with specified dtype and shape.
-
-    Args:
-        value: The value to convert
-        dtype: The target dtype
-        shape: The target shape
-
-    Returns:
-        A numpy array with the specified dtype and shape.
-    """
-    if isinstance(value, np.ndarray):
-        return value.astype(dtype).reshape(shape)
-    else:
-        return np.array(value, dtype=dtype).reshape(shape)