Add Orbax checkpointing guide #2169

SamanehSaadat · 2025-08-20T22:28:56Z

No description provided.

gemini-code-assist

Summary of Changes

Hello @SamanehSaadat, 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 introduces a new guide for Keras users, demonstrating how to implement Orbax checkpointing when training models with the JAX backend. This is particularly beneficial for multi-host training scenarios where the default .keras format may not be suitable.

Highlights

New Orbax Checkpointing Guide: A comprehensive guide (guides/orbax_checkpoint.py) has been added, detailing the process of using Orbax for checkpointing in Keras with the JAX backend.
Custom Checkpointing Utilities: The guide introduces and explains two custom Python classes, KerasOrbaxCheckpointManager and OrbaxCheckpointCallback, which wrap Orbax's functionality for seamless integration with Keras models.
Practical Example: The guide includes a runnable example demonstrating how to set up and use the OrbaxCheckpointCallback during model training.
Guide Integration: The new guide has been registered in scripts/guides_master.py to ensure it appears in the documentation.

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 in your pull request via creating an issue comment (i.e. comment on the pull request page) using either /gemini <command> or @gemini-code-assist <command>. Below is a summary of the supported commands.

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 issue 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.

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 adds a new guide for using Orbax for checkpointing in Keras with the JAX backend. The guide is well-structured and provides useful utility classes. I've provided some suggestions to improve the robustness and usability of the provided OrbaxCheckpointCallback and KerasOrbaxCheckpointManager utilities, including automatically inferring steps_per_epoch, handling potential None values during state restoration, and improving code style and logging practices.

guides/orbax_checkpoint.py

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

… into orbax-guide

hertschuh

Thanks for the PR! This is very useful.

hertschuh · 2025-08-20T22:42:32Z

guides/orbax_checkpoint.py

+"""
+
+"""
+## Introduction


nitpick: add an empty line after each header, like:

## Introduction Orbax is...

guides/orbax_checkpoint.py

hertschuh · 2025-08-20T22:49:16Z

guides/orbax_checkpoint.py

+This guide explains how to do Orbax checkpointing when training a model in
+the Jax backend.
+
+The default `.keras` format doesn't support multi-host checkpointing so if


Let's rephrase this slightly, it's not intrinsically the file format that's wrong. How about:

Note that you should use Orbax checkpointing for multi-host training as the default Keras checkpointing currently does not support multi-host.

Sounds good! Done

hertschuh · 2025-08-20T22:50:01Z

guides/orbax_checkpoint.py

+"""
+
+"""
+We need to set the Keras backend to Jax as this guide is intended for the


I see both JAX and Jax. Let's stay consistent and use JAX everywhere.

hertschuh · 2025-08-20T22:51:29Z

guides/orbax_checkpoint.py

+"""
+
+"""shell
+pip install -q -u orbax-checkpoint


I think it's -U uppercase for upgrade.

Right :D fixed

guides/orbax_checkpoint.py

hertschuh · 2025-08-20T22:55:38Z

guides/orbax_checkpoint.py

+        super().__init__(checkpoint_dir, options=options)
+
+    def _get_state(self):
+        """Gets the model state and metrics"""


Add Returns: to docstring

guides/orbax_checkpoint.py

hertschuh · 2025-08-20T22:56:37Z

guides/orbax_checkpoint.py

+        if keras.config.backend() != "jax":
+            raise ValueError(
+                "`OrbaxCheckpointCallback` is only supported on a "
+                "`jax` backend. Provided backend is %s." % keras.config.backend()


Use f-string here.

hertschuh · 2025-08-20T22:58:34Z

guides/orbax_checkpoint.py

+"""
+Now if you look at the Orbax checkpoint directory, you can see all the files
+saved as part of Orbax checkpointing.
+"""


Oh maybe you can add:

"""shell ls -l /tmp/ckpt """

Does that work?

Orbax checkpoint has nested directories so I initially had this:

"""shell ls -R /tmp/ckpt """

Which outputs:

I wasn't sure if it looks good or not so I removed it! Added it back! Let me know if it looks okay.

hertschuh · 2025-08-20T23:01:17Z

Can you also run it and generate the ipynb and md files?

SamanehSaadat · 2025-08-20T23:57:19Z

Can you also run it and generate the ipynb and md files?

Added ipynb and md files.

hertschuh

Great, thanks!

SamanehSaadat · 2025-08-21T17:01:52Z

Great, thanks!

Thank you for the review!

Add Orbax checkpointing guide

3857036

SamanehSaadat requested review from MarkDaoust and fchollet as code owners August 20, 2025 22:28

gemini-code-assist bot reviewed Aug 20, 2025

View reviewed changes

github-actions bot assigned sachinprasadhs Aug 20, 2025

gemini-code-assist bot reviewed Aug 20, 2025

View reviewed changes

guides/orbax_checkpoint.py Show resolved Hide resolved

guides/orbax_checkpoint.py Show resolved Hide resolved

guides/orbax_checkpoint.py Outdated Show resolved Hide resolved

guides/orbax_checkpoint.py Show resolved Hide resolved

SamanehSaadat and others added 3 commits August 20, 2025 15:41

Make the _get_state function more Pythonic and concise

5d251d2

Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Handle the case that self.latest_step() is None in restore_state

f12e8b1

Merge branch 'orbax-guide' of https://github.com/SamanehSaadat/keras-io…

895a43f

… into orbax-guide

SamanehSaadat requested a review from hertschuh August 20, 2025 22:56

hertschuh reviewed Aug 20, 2025

View reviewed changes

SamanehSaadat added 3 commits August 20, 2025 23:42

Address reviews

6e5a849

Fix some formatting issues

e862207

Add ipynb and md files

688a860

SamanehSaadat requested a review from hertschuh August 21, 2025 15:54

hertschuh approved these changes Aug 21, 2025

View reviewed changes

SamanehSaadat merged commit d3908eb into keras-team:master Aug 21, 2025
3 checks passed

Add Orbax checkpointing guide #2169

Add Orbax checkpointing guide #2169

Uh oh!

Conversation

SamanehSaadat commented Aug 20, 2025

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Summary of Changes

Highlights

Footnotes

Uh oh!

gemini-code-assist bot left a comment

Choose a reason for hiding this comment

Code Review

Uh oh!

Uh oh!

Uh oh!

Uh oh!

Uh oh!

hertschuh left a comment

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

Choose a reason for hiding this comment

Uh oh!

hertschuh commented Aug 20, 2025

Uh oh!

SamanehSaadat commented Aug 20, 2025

Uh oh!

hertschuh left a comment

Choose a reason for hiding this comment

Uh oh!

SamanehSaadat commented Aug 21, 2025

Uh oh!

Uh oh!

Reviewers

Assignees

Labels

Projects

Milestone

Development

Uh oh!

3 participants