Optimize Orbax checkpoint for JAX backend

- Avoid unnecessary numpy conversion in _get_state_tree() for JAX backend
- Preserve JAX arrays during saving instead of converting to numpy
- Maintain cross-backend compatibility with proper loading conversions
- Update async waiting to use CheckpointManager.wait_until_finished()
- Implement AlwaysSavePolicy for reliable save decisions
- Add expected failures for sklearn tests due to neural network non-determinism

Author: amitsrivastava78

keras/src/callbacks/orbax_checkpoint.py

@@ -16,7 +16,12 @@
 
 def _get_state_tree(model):
     """Get the complete model state as a nested tree structure."""
-    state_tree = model.get_state_tree(value_format="numpy_array")
+    # For JAX backend, keep arrays in their native format to avoid unnecessary conversions
+    # For other backends, convert to numpy arrays
+    if backend.backend() == "jax":
+        state_tree = model.get_state_tree()
+    else:
+        state_tree = model.get_state_tree(value_format="numpy_array")
 
     # Convert numpy scalar types to Python types for Orbax compatibility
     def convert_scalars(obj):