fix: reuse the auto-build forward pass in evaluate

synalinks/src/trainers/trainer.py

@@ -831,11 +831,18 @@ async def evaluate(
                 steps_per_execution=self.steps_per_execution,
             )
 
+        # Building calls the program for real on the first batch. That batch is
+        # the first one this evaluation is about to score anyway — the iterator
+        # is unshuffled and rewound below — so the build's predictions are kept
+        # and handed to the first `test_on_batch` instead of being dropped and
+        # recomputed. Nothing is skipped either way; the difference is one
+        # forward pass, which for an agent program is a whole agent run.
+        prebuilt_y_pred = None
         if not all(module.built for module in self._flatten_modules()):
             # Build the model on one batch of data.
             for _, data in epoch_iterator:
                 data_batch = data[0]
-                self._auto_build(
+                prebuilt_y_pred = self._auto_build(
                     iterator=epoch_iterator,
                     data_batch=data_batch,
                 )
@@ -862,10 +869,14 @@ async def evaluate(
             callbacks.on_test_batch_begin(step)
             data = iterator[0]
             x_batch, y_batch = data_adapter_utils.unpack_x_y(data)
+            # Only the first step can consume the auto-build's predictions, and
+            # only once: taking it clears it, so every later batch predicts.
+            y_pred, prebuilt_y_pred = prebuilt_y_pred, None
             logs = await self.test_on_batch(
                 x=x_batch,
                 y=y_batch,
                 return_dict=True,
+                y_pred=y_pred,
             )
             callbacks.on_test_batch_end(step, logs)
             if self.stop_evaluating:
@@ -1024,6 +1035,7 @@ async def test_on_batch(
         x,
         y=None,
         return_dict=False,
+        y_pred=None,
     ):
         """Test the program on a single batch of samples.
 
@@ -1033,14 +1045,22 @@ async def test_on_batch(
             return_dict (bool): If `True`, reward and metric results are returned as a
                 dict, with each key being the name of the metric. If `False`,
                 they are returned as a list.
+            y_pred (list): Optional predictions for `x`, already computed by an
+                earlier forward pass. When given, the program is not called
+                again — see `evaluate`, which reuses the prediction its
+                auto-build pass produced for the first batch. Passing
+                predictions that are not the program's own output for `x` will
+                silently report rewards and metrics for something the program
+                never predicted.
 
         Returns:
             (float | list | dict): A scalar reward value
                 (when no metrics and `return_dict=False`), a list of reward
                 and metric values (if there are metrics and `return_dict=False`),
                 or a dict of metric and reward values (if `return_dict=True`).
         """
-        y_pred = await self.predict_on_batch(x)
+        if y_pred is None:
+            y_pred = await self.predict_on_batch(x)
 
         rewards = await self.compute_reward(
             x=x,
@@ -1193,6 +1213,20 @@ def _assert_compile_called(self, method_name=None):
             raise ValueError(msg)
 
     def _auto_build(self, iterator=None, data_batch=None):
+        """Build the program, its metrics, its reward and its optimizer.
+
+        Building the program means calling it: a program's output schema is
+        not knowable without running it, so this does one real forward pass on
+        `data_batch`. That pass is as expensive as any other — for an agent
+        program it is a full agent run, LM calls and all — so the predictions
+        it produces are returned rather than dropped, letting a caller that is
+        about to evaluate the very same batch reuse them.
+
+        Returns:
+            (list | None): The predictions for `data_batch`'s inputs, or None
+                when nothing needed building and no forward pass was made.
+        """
+        y_pred = None
         program_unbuilt = not all(module.built for module in self._flatten_modules())
         compile_metrics_unbuilt = (
             self._compile_metrics is not None and not self._compile_metrics.built
@@ -1249,6 +1283,7 @@ def _auto_build(self, iterator=None, data_batch=None):
             # Build optimizer
             run_maybe_nested(self.optimizer.build(self.trainable_variables))
         self._post_build()
+        return y_pred
 
     def _assert_compile_called(self, method_name=None):
         if not self.compiled:

synalinks/src/trainers/trainer_test.py

@@ -145,6 +145,91 @@ async def test_evaluate(self, mock_completion):
             y=y_test,
         )
 
+    @patch("litellm.acompletion")
+    async def test_evaluate_predicts_each_sample_once_when_unbuilt(
+        self, mock_completion
+    ):
+        """An unbuilt program must not be run twice on the batch it builds on.
+
+        Building a program means calling it, and `evaluate` builds on the first
+        batch — which is also the first batch it scores. Dropping that build's
+        predictions and recomputing them costs a second full forward pass per
+        sample in the batch, which for an agent program is a second agent run:
+        every LM call and every tool call paid twice, for a result thrown away.
+
+        The program here keeps a module its `call()` never reaches, which is
+        what makes this permanent rather than a one-off. `built` is only set by
+        calling a module, so a module on an untaken path never becomes built,
+        `all(module.built ...)` never becomes true, and *every* `evaluate` pays
+        the extra pass. Agents are full of such modules — a fallback generator
+        that only runs when the loop ends without submitting, a tool the model
+        happens not to call.
+        """
+        mock_answer = AnswerWithRationale(
+            rationale="""The capital of France is well-known and is the seat of """
+            """the French government.""",
+            answer="Paris",
+        )
+
+        mock_completion.return_value = {
+            "choices": [{"message": {"content": json.dumps(mock_answer.get_json())}}]
+        }
+
+        language_model = LanguageModel(model="ollama/mistral")
+
+        class ProgramWithAnUntakenPath(programs.Program):
+            forward_passes = 0
+
+            def __init__(self, language_model=None):
+                super().__init__()
+                self.answer = modules.Generator(
+                    data_model=AnswerWithRationale,
+                    language_model=language_model,
+                )
+                self.fallback = modules.Generator(
+                    data_model=AnswerWithRationale,
+                    language_model=language_model,
+                )
+
+            async def call(self, inputs, training=False):
+                type(self).forward_passes += 1
+                return await self.answer(inputs)
+
+        program = ProgramWithAnUntakenPath(language_model=language_model)
+
+        program.compile(
+            optimizer=optimizers.random_few_shot.RandomFewShot(),
+            reward=rewards.ExactMatch(in_mask=["answer"]),
+            metrics=[
+                metrics.MeanMetricWrapper(rewards.exact_match, in_mask=["answer"]),
+            ],
+        )
+
+        (_, _), (x_test, y_test) = load_test_data()
+
+        self.assertFalse(all(module.built for module in program._flatten_modules()))
+
+        _ = await program.evaluate(x=x_test, y=y_test)
+
+        # The build pass predicted the whole first batch, and those predictions
+        # are what the first `test_on_batch` scores — so each sample is
+        # predicted once, not twice. The one extra pass is `_auto_build`'s own
+        # spec pass for the metric and reward state; it is paid once per
+        # program, not once per sample, and is not what this reuse addresses.
+        self.assertLessEqual(
+            ProgramWithAnUntakenPath.forward_passes, len(x_test) + 1
+        )
+
+        # `fallback` is still unbuilt, so the next `evaluate` auto-builds again
+        # — and must again cost exactly one pass per sample. This is what makes
+        # the duplication permanent rather than a first-call toll.
+        before = ProgramWithAnUntakenPath.forward_passes
+        _ = await program.evaluate(x=x_test, y=y_test)
+        self.assertFalse(all(module.built for module in program._flatten_modules()))
+        self.assertEqual(
+            ProgramWithAnUntakenPath.forward_passes - before, len(x_test)
+        )
+
     @patch("litellm.acompletion")
     async def test_predict(self, mock_completion):
         mock_answer = AnswerWithRationale(