Add OrbaxCheckpoint callback with conditional exports and improved test handling

- Implement OrbaxCheckpoint callback for async checkpointing with state tree handling
- Add conditional exports for optional orbax-checkpoint dependency
- Use pytest.importorskip for clean optional dependency testing
- Ensure graceful handling when orbax-checkpoint is not installed

Author: amitsrivastava78

keras/src/callbacks/orbax_checkpoint.py

@@ -21,21 +21,9 @@
     ),
 )
 
-# Expose advanced Orbax functionality for users who need direct access
-# These are provided as bridge for advanced usecases like custom type handlers
-CheckpointManager = ocp.CheckpointManager
-SaveArgs = ocp.SaveArgs
-StandardRestore = ocp.args.StandardRestore
-
-# Type handler functionality for custom serialization
-TypeHandler = ocp.type_handlers.TypeHandler
-register_type_handler = ocp.type_handlers.register_type_handler
-
-# Direct checkpointing for custom objects
-PyTreeCheckpointer = ocp.PyTreeCheckpointer
-
-# Metadata functionality
-metadata = ocp.metadata
+# Note: Advanced Orbax functionality is available through the ocp LazyModule
+# Users can access it via: from keras.src.utils.module_utils import LazyModule
+# ocp = LazyModule("orbax.checkpoint"); ocp.CheckpointManager
 
 
 def _get_state_tree(model):
@@ -701,3 +689,14 @@ def _restore_from_state_tree(self, state_tree, target_model):
         if self.verbose > 0:
             print_msg("OrbaxCheckpoint: Successfully restored model state")
         return True
+
+
+# Export additional Orbax functionality for advanced users (only if available)
+if ocp.available:
+    CheckpointManager = ocp.CheckpointManager
+    PyTreeCheckpointer = ocp.PyTreeCheckpointer
+    SaveArgs = ocp.SaveArgs
+    StandardRestore = ocp.args.StandardRestore
+    TypeHandler = ocp.type_handlers.TypeHandler
+    metadata = ocp.metadata
+    register_type_handler = ocp.type_handlers.register_type_handler

keras/src/callbacks/orbax_checkpoint_test.py

@@ -11,14 +11,23 @@
 from keras.src import testing
 
 # Import advanced Orbax functionality through the Keras bridge
-from keras.src.callbacks.orbax_checkpoint import CheckpointManager
+# These will only be available if orbax-checkpoint is installed
+try:
+    from keras.src.callbacks.orbax_checkpoint import CheckpointManager
+    from keras.src.callbacks.orbax_checkpoint import PyTreeCheckpointer
+    from keras.src.callbacks.orbax_checkpoint import SaveArgs
+    from keras.src.callbacks.orbax_checkpoint import StandardRestore
+    from keras.src.callbacks.orbax_checkpoint import TypeHandler
+    from keras.src.callbacks.orbax_checkpoint import metadata
+    from keras.src.callbacks.orbax_checkpoint import register_type_handler
+except ImportError:
+    # If orbax is not available, these won't be exported
+    pass
+
 from keras.src.callbacks.orbax_checkpoint import OrbaxCheckpoint
-from keras.src.callbacks.orbax_checkpoint import PyTreeCheckpointer
-from keras.src.callbacks.orbax_checkpoint import SaveArgs
-from keras.src.callbacks.orbax_checkpoint import StandardRestore
-from keras.src.callbacks.orbax_checkpoint import TypeHandler
-from keras.src.callbacks.orbax_checkpoint import metadata
-from keras.src.callbacks.orbax_checkpoint import register_type_handler
+
+# Skip the entire test module if orbax-checkpoint is not available
+pytest.importorskip("orbax.checkpoint")
 
 
 class OrbaxCheckpointTest(testing.TestCase):