Implement complete Keras-Orbax checkpoint integration

[amitsrivastava78]

• Add compile config saving to OrbaxCheckpoint for full model restoration
• Update saving_api.py to use compile_from_config for proper model loading
• Replace direct Orbax API tests with Keras saving API tests
• Remove unused Orbax imports from test file
• Add comprehensive test coverage for all model state components
• Ensure cross-backend compatibility (TensorFlow, JAX, PyTorch)

This enables users to save and load complete model state (weights, optimizer, metrics) using keras.saving.load_model() with Orbax checkpoints.

[gemini-code-assist]

Summary of Changes

Hello @amitsrivastava78, I'm Gemini Code Assist¹! I'm currently reviewing this pull request and will post my feedback shortly. In the meantime, here's a summary to help you and other reviewers quickly get up to speed!

This pull request significantly enhances the integration of Orbax checkpointing within Keras by enabling the saving and loading of a complete model state, including its architecture, compilation settings, optimizer, and metrics. This allows users to fully restore a trained model from an Orbax checkpoint using the standard keras.saving.load_model() function, eliminating the need for manual reconstruction or separate loading steps for different model components. The changes ensure a more robust and user-friendly experience for checkpoint management.

Highlights

• Complete Model State Saving: The OrbaxCheckpoint callback now saves the full model configuration, compile configuration, and optimizer configuration, enabling comprehensive model restoration.
• Unified Loading API: The keras.saving.load_model() function has been extended to directly support loading models from Orbax checkpoint directories, streamlining the workflow for users.
• Enhanced Test Coverage: Existing tests for Orbax checkpoint loading have been refactored to use the Keras saving API, and new tests have been added to ensure complete model state (weights, optimizer, metrics, compile config) is correctly saved and loaded across different backends.
• Cross-Backend Compatibility: The loading mechanism includes logic to convert JAX arrays to NumPy arrays for non-JAX backends, ensuring seamless operation regardless of the underlying Keras backend.

🧠 New Feature in Public Preview: You can now enable Memory to help Gemini Code Assist learn from your team's feedback. This makes future code reviews more consistent and personalized to your project's style. Click here to enable Memory in your admin console.

Using Gemini Code Assist

The full guide for Gemini Code Assist can be found on our documentation page, here are some quick tips.

Invoking Gemini

You can request assistance from Gemini at any point by creating a comment using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands on the current page.

Feature	Command	Description
Code Review	/gemini review	Performs a code review for the current pull request in its current state.
Pull Request Summary	/gemini summary	Provides a summary of the current pull request in its current state.
Comment	@gemini-code-assist	Responds in comments when explicitly tagged, both in pull request comments and review comments.
Help	/gemini help	Displays a list of available commands.

Customization

To customize Gemini Code Assist for GitHub experience, repository maintainers can create a configuration file and/or provide a custom code review style guide (such as PEP-8 for Python) by creating and adding files to a .gemini/ folder in the base of the repository. Detailed instructions can be found here.

Limitations & Feedback

Gemini Code Assist may make mistakes. Please leave feedback on any instances where its feedback is incorrect or counter productive. You can react with 👍 and 👎 on @gemini-code-assist comments. If you're interested in giving your feedback about your experience with Gemini Code Assist for Github and other Google products, sign up here.

You can also get AI-powered code generation, chat, as well as code reviews directly in the IDE at no cost with the Gemini Code Assist IDE Extension.

Footnotes

1. Review the Privacy Notices, Generative AI Prohibited Use Policy, Terms of Service, and learn how to configure Gemini Code Assist in GitHub here. Gemini can make mistakes, so double check it and use code with caution. ↩

[gemini-code-assist]

Code Review

This pull request does a great job of integrating Keras with Orbax for checkpointing, enabling full model restoration via keras.saving.load_model. The changes to OrbaxCheckpoint to save model and compile configurations are correct. The new loading logic in saving_api.py is robust, handling detection of Orbax checkpoints and reconstruction of the model state. The test suite has been significantly improved by replacing direct Orbax API calls with tests against the public Keras saving API, and by adding comprehensive tests for all components of the model state.

I have a couple of suggestions to improve the code further. One is a high-severity suggestion to ensure remote path support in the new loading logic, and the other is a medium-severity suggestion to refactor duplicated code in the tests for better maintainability. Overall, this is a solid contribution.

[gemini-code-assist]

For consistency with the rest of Keras and to support remote file systems (like GCS), you should use the wrappers from keras.src.utils.file_utils (e.g., file_utils.exists, file_utils.listdir, file_utils.isdir) instead of the os module. The detection logic in load_model already uses file_utils, and this function should too to ensure it works correctly with remote paths.

Suggested change

	if os.path.exists(filepath):
	subdirs = os.listdir(filepath)
	for d in subdirs:
	if os.path.isdir(os.path.join(filepath, d)):
	try:
	available_steps.append(int(d))
	except ValueError:
	pass
	if file_utils.exists(filepath):
	subdirs = file_utils.listdir(filepath)
	for d in subdirs:
	if file_utils.isdir(file_utils.join(filepath, d)):
	try:
	available_steps.append(int(d))
	except ValueError:
	pass

[codecov-commenter]

Codecov Report

❌ Patch coverage is 90.62500% with 3 lines in your changes missing coverage. Please review.
✅ Project coverage is 82.80%. Comparing base (fdc5543) to head (62b8cb6).
⚠️ Report is 254 commits behind head on master.

Files with missing lines	Patch %	Lines
keras/src/callbacks/orbax_checkpoint.py	80.00%	3 Missing ⚠️

Additional details and impacted files

@@            Coverage Diff             @@
##           master   #22002      +/-   ##
==========================================
+ Coverage   82.69%   82.80%   +0.10%     
==========================================
  Files         589      592       +3     
  Lines       61632    62543     +911     
  Branches     9650     9797     +147     
==========================================
+ Hits        50967    51787     +820     
- Misses       8165     8220      +55     
- Partials     2500     2536      +36     

Flag	Coverage Δ
keras	82.62% <90.62%> (+0.10%)	⬆️
keras-jax	62.39% <90.62%> (+0.92%)	⬆️
keras-numpy	56.45% <12.50%> (-0.27%)	⬇️
keras-openvino	37.61% <12.50%> (+0.13%)	⬆️
keras-tensorflow	63.66% <90.62%> (+0.02%)	⬆️
keras-torch	62.42% <90.62%> (+0.03%)	⬆️

Flags with carried forward coverage won't be shown. Click here to find out more.

☔ View full report in Codecov by Sentry.
📢 Have feedback on the report? Share it here.

🚀 New features to boost your workflow:

• ❄️ Test Analytics: Detect flaky tests, report on failures, and find test suite problems.

[hertschuh]

The whole asset saving / loading part of it is missing.

[amitsrivastava78]

The whole asset saving / loading part of it is missing.

yes will raise a separate PR for that

[hertschuh]

I realized that the config saving and loading should reuse some utilities that already exist and take care of some non-trivial stuff.

[amitsrivastava78]

I realized that the config saving and loading should reuse some utilities that already exist and take care of some non-trivial stuff.

yes, the implementation now is perfectly aligned with the recommendation to reuse existing utilities

[hertschuh]

Thanks!