orbit: fix LoopFnWithSummaries on zero steps, add regression test

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(

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

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):