Optimize Orbax checkpoint for JAX backend with compatibility check

- Preserve JAX arrays during saving when jax.monitoring.record_scalar is available
- Fall back to numpy conversion for older JAX versions that don't have record_scalar
- Maintain cross-backend compatibility while avoiding unnecessary 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,10 +16,17 @@
 
 def _get_state_tree(model):
     """Get the complete model state as a nested tree structure."""
-    # For JAX backend, preserve native arrays to avoid unnecessary conversions
-    # For other backends, convert to numpy arrays
+    # For JAX backend, preserve native arrays if JAX monitoring is available
+    # to avoid unnecessary conversions. Otherwise convert to numpy for compatibility.
     if backend.backend() == "jax":
-        state_tree = model.get_state_tree()
+        try:
+            import jax
+            # Check if jax.monitoring.record_scalar exists (introduced in JAX 0.7.0)
+            jax.monitoring.record_scalar
+            state_tree = model.get_state_tree()
+        except (ImportError, AttributeError):
+            # Fallback to numpy conversion for older JAX versions
+            state_tree = model.get_state_tree(value_format="numpy_array")
     else:
         state_tree = model.get_state_tree(value_format="numpy_array")