Improve OrbaxCheckpoint implementation

- Remove conditional export decorator to ensure OrbaxCheckpoint is always available
- Remove unnecessary exception handling in state tree operations
- Update process index check comment for clarity
- Format code to comply with 80-character line limit
- Add distribution_lib modules for backend-specific distributed training support

Author: amitsrivastava78

keras/src/backend/jax/__init__.py

@@ -1,6 +1,5 @@
 from keras.src.backend.config import is_nnx_enabled
 from keras.src.backend.jax import core
-from keras.src.backend.jax import distribution_lib
 from keras.src.backend.jax import image
 from keras.src.backend.jax import linalg
 from keras.src.backend.jax import math
@@ -29,3 +28,4 @@
 from keras.src.backend.jax.rnn import gru
 from keras.src.backend.jax.rnn import lstm
 from keras.src.backend.jax.rnn import rnn
+from keras.src.backend.jax.distribution_lib import process_id

keras/src/backend/numpy/__init__.py

@@ -24,3 +24,4 @@
 from keras.src.backend.numpy.rnn import gru
 from keras.src.backend.numpy.rnn import lstm
 from keras.src.backend.numpy.rnn import rnn
+from keras.src.backend.numpy.distribution_lib import process_id

keras/src/backend/numpy/distribution_lib.py

@@ -0,0 +1,6 @@
+"""Utilities for distribution strategy with NumPy backend."""
+
+
+def process_id():
+    """Return the current process ID for the distribution setting."""
+    return 0

keras/src/backend/openvino/__init__.py

@@ -1,5 +1,6 @@
 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
@@ -23,3 +24,4 @@
 from keras.src.backend.openvino.rnn import gru
 from keras.src.backend.openvino.rnn import lstm
 from keras.src.backend.openvino.rnn import rnn
+from keras.src.backend.openvino.distribution_lib import process_id

keras/src/backend/openvino/distribution_lib.py

@@ -0,0 +1,6 @@
+"""Utilities for distribution strategy with OpenVINO backend."""
+
+
+def process_id():
+    """Return the current process ID for the distribution setting."""
+    return 0

keras/src/backend/tensorflow/__init__.py

@@ -1,5 +1,4 @@
 from keras.src.backend.tensorflow import core
-from keras.src.backend.tensorflow import distribution_lib
 from keras.src.backend.tensorflow import image
 from keras.src.backend.tensorflow import linalg
 from keras.src.backend.tensorflow import math
@@ -28,3 +27,4 @@
 from keras.src.backend.tensorflow.rnn import gru
 from keras.src.backend.tensorflow.rnn import lstm
 from keras.src.backend.tensorflow.rnn import rnn
+from keras.src.backend.tensorflow.distribution_lib import process_id

keras/src/backend/tensorflow/distribution_lib.py

@@ -85,3 +85,13 @@ def _to_backend_layout(tensor_layout):
     ]
     dtensor_mesh = tensor_layout.device_mesh.backend_mesh
     return dtensor.Layout(sharding_specs=sharding_specs, mesh=dtensor_mesh)
+
+
+def process_id():
+    """Return the current process ID for the distribution setting."""
+    try:
+        import tensorflow as tf
+
+        return tf.distribute.get_replica_context().replica_id_in_sync_group
+    except (ImportError, AttributeError, RuntimeError):
+        return 0

keras/src/backend/torch/__init__.py

@@ -43,3 +43,4 @@
 from keras.src.backend.torch.rnn import gru
 from keras.src.backend.torch.rnn import lstm
 from keras.src.backend.torch.rnn import rnn
+from keras.src.backend.torch.distribution_lib import process_id

keras/src/backend/torch/distribution_lib.py

@@ -0,0 +1,13 @@
+"""Utilities for distribution strategy with PyTorch backend."""
+
+
+def process_id():
+    """Return the current process ID for the distribution setting."""
+    try:
+        import torch.distributed as dist
+
+        if dist.is_available() and dist.is_initialized():
+            return dist.get_rank()
+        return 0
+    except (ImportError, AttributeError):
+        return 0

keras/src/callbacks/__init__.py

@@ -8,12 +8,7 @@
 from keras.src.callbacks.learning_rate_scheduler import LearningRateScheduler
 from keras.src.callbacks.model_checkpoint import ModelCheckpoint
 from keras.src.callbacks.monitor_callback import MonitorCallback
-
-try:
-    from keras.src.callbacks.orbax_checkpoint import OrbaxCheckpoint
-except ImportError:
-    OrbaxCheckpoint = None
-
+from keras.src.callbacks.orbax_checkpoint import OrbaxCheckpoint
 from keras.src.callbacks.progbar_logger import ProgbarLogger
 from keras.src.callbacks.reduce_lr_on_plateau import ReduceLROnPlateau
 from keras.src.callbacks.remote_monitor import RemoteMonitor