Repository navigation
orbit: fix LoopFnWithSummaries on zero steps, add regression test - #13664
Open
rootkiller6788 wants to merge 3 commits into
Open
rootkiller6788 wants to merge 3 commits into
rootkiller6788 wants to merge 3 commits into
Conversation
LoopFnWithSummaries.__call__ only assigns output in one of its two branches, so calling it with num_steps == 0 threw UnboundLocalError when summaries were off. With summaries on it did something worse: it ran one step and bumped global_step even though zero were requested. create_loop_fn already treats zero steps as a no-op, so make this one agree. Also drops the pyrefly suppression that was hiding the unbound name instead of fixing it.
TrainFunctionWithSummaries in tpu_summaries_test.py is a copy of the LoopFnWithSummaries pattern and had the same hole (and the same suppression). Keep the two in sync.
Goes through StandardTrainer so it covers the real path: create_train_loop_fn wraps the while-loop fn in LoopFnWithSummaries when use_tpu_summary_optimization is on.
rootkiller6788
marked this pull request as ready for review
October 9, 2026 15:19
This file contains hidden or bidirectional Unicode text that may be interpreted or compiled differently than what appears below. To review, open the file in an editor that reveals hidden Unicode characters.
Learn more about bidirectional Unicode characters
Sign up for free
to join this conversation on GitHub.
Already have an account?
Sign in to comment
1 participant
Add this suggestion to a batch that can be applied as a single commit.This suggestion is invalid because no changes were made to the code.Suggestions cannot be applied while the pull request is closed.Suggestions cannot be applied while viewing a subset of changes.Only one suggestion per line can be applied in a batch.Add this suggestion to a batch that can be applied as a single commit.Applying suggestions on deleted lines is not supported.You must change the existing code in this line in order to create a valid suggestion.Outdated suggestions cannot be applied.This suggestion has been applied or marked resolved.Suggestions cannot be applied from pending reviews.Suggestions cannot be applied on multi-line comments.Suggestions cannot be applied while the pull request is queued to merge.Suggestion cannot be applied right now. Please check back later.
LoopFnWithSummaries.call only assigns
outputin one of its two branches, so training with num_steps=0 was broken two different ways: with summary recording off it raised UnboundLocalError, and with recording on it quietly ran one extra step and bumped global_step even though zero were asked for. create_loop_fn already treats zero steps as a no-op, so I made this variant agree and dropped the pyrefly suppression that was papering over it. Same fix in the copy of this pattern in tpu_summaries_test.py.I hit this while reading the pyrefly suppressions added in 8b12ae2 - the unbound-name one there is a real crash, not a false positive.
There's no test for the tpu summary optimization path with zero steps, so I added one to standard_runner_test.py. It goes through StandardTrainer, which is how this actually gets used: create_train_loop_fn wraps the while-loop fn in LoopFnWithSummaries when use_tpu_summary_optimization is on.
I could not run the orbit tests here (no TF on this machine, Python 3.14), so I checked the control flow directly: loaded orbit/utils/loop_fns.py with tensorflow stubbed out and called call for recording on/off x num_steps 0/1/3. Before the change num_steps=0 gave UnboundLocalError (recording off) or a spurious call [('with', 1)] (recording on); after it is [] for both, and num_steps=1 and 3 behave exactly as before. Worth a run on a machine with TF.