Fix checkpoint saving using wrong interval config

[Mr-Neutr0n]

Summary

The checkpoint saving condition checked config.checkpoint_every to determine whether checkpointing was enabled, but used config.eval_every for the modulo operation. That caused checkpoints to follow evaluation intervals instead of the configured checkpoint interval.

The fix now routes the decision through a small _should_checkpoint helper that consistently uses config.checkpoint_every. The final training step is still always checkpointed, including when periodic checkpointing is disabled.

Regression coverage

The previous integration test used total_steps = 1, so it always exercised the final-step fallback and could not catch the interval mix-up. The new parameterized cases cover:

• a configured checkpoint interval
• a step between checkpoint intervals while eval_every is due
• disabled periodic checkpointing
• the final-step fallback when periodic checkpointing is disabled

Verification

• Python compilation passed.
• Four focused regression cases passed on Python 3.10.
• The directory-backed end-to-end training case passed.
• The complete train_test module was attempted; its existing TFDS variant failed locally because the resolved tensorflow_datasets module lacks the expected builder API. That failure is unrelated to this change.

---

This change was prepared with AI assistance under human direction and review.

[google-cla]

Thanks for your pull request! It looks like this may be your first contribution to a Google open source project. Before we can look at your pull request, you'll need to sign a Contributor License Agreement (CLA).

View this failed invocation of the CLA check for more information.

For the most up to date status, view the checks section at the bottom of the pull request.

[Mr-Neutr0n]

Hi! Quick status: this still applies cleanly to the latest main as of today, and the CI/mergeable state is mergeable with all current checks passing. Re-requesting a review — happy to update for any feedback, and just as happy to close if the repo has moved on.