Features: Add sharding and multi-host support for JAX backend

- Sharding support: Enable distributed arrays across JAX devices
- Multi-host support: Coordinate checkpointing across multiple processes
- Interoperability: Load sharded checkpoints to unsharded models and vice versa
- Error handling: Proper validation and backend-specific restrictions
- Comprehensive testing: 11 new test methods covering all scenarios

Author: amitsrivastava78

keras/src/callbacks/orbax_checkpoint.py

@@ -189,6 +189,15 @@ class OrbaxCheckpoint(MonitorCallback):
 
     model.fit(epochs=EPOCHS, callbacks=[orbax_checkpoint_callback])
 
+    # Or use a SaveDecisionPolicy for more control -
+    from orbax.checkpoint import checkpoint_managers
+    policy = checkpoint_managers.FixedIntervalPolicy(interval=5)
+    orbax_checkpoint_callback = keras.callbacks.OrbaxCheckpoint(
+        directory=checkpoint_dir,
+        save_decision_policy=policy)  # Save every 5 epochs
+
+    model.fit(epochs=EPOCHS, callbacks=[orbax_checkpoint_callback])
+
     # JAX-specific features: Sharding and Multi-Host Checkpointing
     # Note: These features are only available with JAX backend
 
@@ -336,7 +345,6 @@ def __init__(
                     "sharding and multi_host parameters are only supported "
                     "with JAX backend. Current backend: " + backend.backend()
                 )
-
         self._batches_seen_since_last_saving = 0
         self._last_batch_seen = 0
         self._current_epoch = 0  # Keep track of epoch
@@ -506,10 +514,7 @@ def _save_checkpoint(self, step, logs=None):
                 composite_state["data_iterator"] = iterator_state
 
         # --- Save Logic ---
-        # In multi-host setups, only the primary process (rank 0) initiates the
-        # save operation. Orbax internally coordinates distributed writing: each
-        # process writes its own data shards in parallel while the primary
-        # process manages metadata and coordination.
+        # Only save on the primary process (rank 0) in distributed setups
         is_primary_host = backend.get_process_index() == 0
 
         if is_primary_host: