Fix OrbaxCheckpoint sharding and multi-host issues

- Fix sharding parameter passing in save/restore operations by passing as kwargs instead of setting attributes on StandardSave/StandardRestore objects
- Add robust error handling for distribution initialization with multiple error message patterns
- Add proper test skipping for JAX-only features when distribution module unavailable
- Add sharding parameter validation in constructor to prevent invalid types
- Update test expectations to match corrected sharding validation behavior

These changes ensure proper sharding support for JAX multi-host checkpointing while maintaining backward compatibility.

Author: amitsrivastava78

keras/src/callbacks/orbax_checkpoint.py

@@ -345,6 +345,16 @@ def __init__(
                     "sharding and multi_host parameters are only supported "
                     "with JAX backend. Current backend: " + backend.backend()
                 )
+
+        # Validate sharding object type
+        if sharding is not None and backend.backend() == "jax":
+            # Basic validation: sharding should not be a string or other
+            # primitive type
+            if isinstance(sharding, (str, int, float, bool)):
+                raise TypeError(
+                    f"sharding parameter must be a valid JAX sharding object, "
+                    f"got {type(sharding).__name__}: {sharding}"
+                )
         self._batches_seen_since_last_saving = 0
         self._last_batch_seen = 0
         self._current_epoch = 0  # Keep track of epoch
@@ -395,9 +405,14 @@ def __init__(
             except RuntimeError as e:
                 # If distributed cannot be initialized (e.g., JAX already
                 # initialized), continue anyway - the multi_host flag is mainly
-                # a hint to Orbax
-                if "must be called before" in str(e):
-                    pass  # This is expected in test environments
+                # a hint to Orbax.
+                # We check for messages related to initialization state.
+                error_str = str(e).lower()
+                if (
+                    "already been initialized" in error_str
+                    or "must be called before" in error_str
+                ):
+                    pass  # This is expected in some environments.
                 else:
                     raise
             # Orbax will automatically handle multi-host coordination:
@@ -529,15 +544,10 @@ def _save_checkpoint(self, step, logs=None):
             )
 
             # Apply sharding if specified (JAX only)
+            save_kwargs = {}
             if self.sharding is not None and backend.backend() == "jax":
-                # For JAX sharding, we need to ensure the data is properly
-                # sharded
-                # This is typically handled automatically by Orbax when JAX
-                # arrays with sharding metadata are saved
-                if hasattr(save_args, "sharding"):
-                    save_args.sharding = self.sharding
-
-            self.manager.save(step, args=save_args)
+                save_kwargs["sharding"] = self.sharding
+            self.manager.save(step, args=save_args, **save_kwargs)
 
     def on_train_batch_end(self, batch, logs=None):
         if self._should_save_on_batch(batch):
@@ -650,14 +660,14 @@ def load_checkpoint(self, step, model=None):
         restore_args = ocp.args.StandardRestore()
 
         # Apply sharding if specified (JAX only)
+        restore_kwargs = {}
         if self.sharding is not None and backend.backend() == "jax":
-            # For JAX sharding, we need to ensure the data is properly restored
-            # with the same sharding specification used during save
-            if hasattr(restore_args, "sharding"):
-                restore_args.sharding = self.sharding
+            restore_kwargs["sharding"] = self.sharding
 
         # Load the checkpoint
-        checkpoint_data = self.manager.restore(step, args=restore_args)
+        checkpoint_data = self.manager.restore(
+            step, args=restore_args, **restore_kwargs
+        )
 
         # Restore the model state
         target_model = model if model is not None else self.model

keras/src/callbacks/orbax_checkpoint_test.py

@@ -2044,10 +2044,14 @@ def test_restore_unsharded_checkpoint_to_sharded_model(self):
                 "unsharded checkpoint",
             )
 
+    @pytest.mark.skipif(
+        backend.backend() != "jax",
+        reason="Sharding validation tests require JAX backend",
+    )
     def test_invalid_sharding_argument_raises_error(self):
         """Test that invalid sharding arguments raise TypeError."""
         # Test with string (invalid sharding object)
-        with self.assertRaises((TypeError, ValueError)):
+        with self.assertRaises(TypeError):
             OrbaxCheckpoint(
                 directory=os.path.join(self.temp_dir, "test_invalid_sharding"),
                 sharding="invalid_sharding_string",