From aa3b682af32172d65f937cf24eb060233b3276b4 Mon Sep 17 00:00:00 2001
From: rootkiller6788 <221446036+rootkiller6788@users.noreply.github.com>
Date: Thu, 17 Sep 2026 13:52:03 +0800
Subject: [PATCH 1/3] orbit: keep LoopFnWithSummaries from blowing up on zero
 steps

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.
---
 orbit/utils/loop_fns.py | 5 ++++-
 1 file changed, 4 insertions(+), 1 deletion(-)

diff --git a/orbit/utils/loop_fns.py b/orbit/utils/loop_fns.py
index 2cbcc6fa25e..1eaac041a91 100644
--- a/orbit/utils/loop_fns.py
+++ b/orbit/utils/loop_fns.py
@@ -203,10 +203,13 @@ class LoopFnWithSummaries(tpu_summaries.OptionalSummariesFunction):
   """
 
   def __call__(self, iterator, num_steps):
+    if num_steps < 1:
+      # Zero steps is a no-op, same as in `create_loop_fn`.
+      return
     if tf.summary.should_record_summaries():
       with ecu.SequentialEmbeddingContext():
         output = self.with_summaries(iterator, tf.constant(1))
       num_steps -= 1
     if num_steps >= 1:
       output = self.without_summaries(iterator, num_steps)
-    return output  # pyrefly: ignore[unbound-name]
+    return output

From 287790cc6397068a8c77ea64cc3278172edb46d4 Mon Sep 17 00:00:00 2001
From: rootkiller6788 <221446036+rootkiller6788@users.noreply.github.com>
Date: Thu, 17 Sep 2026 13:52:09 +0800
Subject: [PATCH 2/3] orbit: same zero-step handling in the tpu_summaries test
 helper

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.
---
 orbit/utils/tpu_summaries_test.py | 4 +++-
 1 file changed, 3 insertions(+), 1 deletion(-)

diff --git a/orbit/utils/tpu_summaries_test.py b/orbit/utils/tpu_summaries_test.py
index c1194007ed8..8057f8ed662 100644
--- a/orbit/utils/tpu_summaries_test.py
+++ b/orbit/utils/tpu_summaries_test.py
@@ -27,12 +27,14 @@ class TrainFunctionWithSummaries(tpu_summaries.OptionalSummariesFunction):
   """Implements a two-program approach for summaries on TPU."""
 
   def __call__(self, num_steps):
+    if num_steps < 1:
+      return
     if tf.summary.should_record_summaries():
       output = self.with_summaries(tf.constant(1))
       num_steps -= 1
     if num_steps >= 1:
       output = self.without_summaries(num_steps)
-    return output  # pyrefly: ignore[unbound-name]
+    return output
 
 
 def train_function_with_summaries(function=None, **kwargs):

From 76f56210e644ef05e7656431380315a2b0e02bf2 Mon Sep 17 00:00:00 2001
From: rootkiller6788 <221446036+rootkiller6788@users.noreply.github.com>
Date: Thu, 17 Sep 2026 13:52:09 +0800
Subject: [PATCH 3/3] orbit: regression test for training zero steps with tpu
 summary opt

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.
---
 orbit/standard_runner_test.py | 6 ++++++
 1 file changed, 6 insertions(+)

diff --git a/orbit/standard_runner_test.py b/orbit/standard_runner_test.py
index 57663fe23c3..85acdeb66e3 100644
--- a/orbit/standard_runner_test.py
+++ b/orbit/standard_runner_test.py
@@ -122,6 +122,12 @@ def test_trainer_with_tpu_summary_optimization(self):
     trainer = TestTrainer(options)
     self.assertEqual(trainer.train(tf.constant(10)), 10)
 
+  def test_trainer_with_tpu_summary_optimization_zero_steps(self):
+    options = standard_runner.StandardTrainerOptions(
+        use_tpu_summary_optimization=True)
+    trainer = TestTrainer(options)
+    self.assertEqual(trainer.train(tf.constant(0)), 0)
+
   @parameterized.named_parameters(("use_tf_while_loop", True), ("", False))
   def test_default_evaluator(self, use_tf_while_loop):
     options = standard_runner.StandardEvaluatorOptions(