Add return_context to trainstep.step()

PiperOrigin-RevId: 699904187

kauldron/evals/evaluators.py

@@ -251,7 +251,6 @@ def evaluate(
         step=step,
         aux=merged_aux,
         schedules={},
-        model_with_aux=self.model_with_aux,
         log_summaries=True,
     )
     return merged_aux

kauldron/train/__init__.py

@@ -21,7 +21,9 @@
 from kauldron.train.setup_utils import Setup
 from kauldron.train.setup_utils import TqdmInfo
 from kauldron.train.train_step import Auxiliaries
-from kauldron.train.train_step import ModelWithAux
+from kauldron.train.train_step import AuxiliariesRef
+from kauldron.train.train_step import forward
+from kauldron.train.train_step import forward_with_loss
 from kauldron.train.train_step import TrainState
 from kauldron.train.train_step import TrainStep
 from kauldron.train.trainer_lib import Trainer

kauldron/train/context.py

@@ -62,6 +62,8 @@ class Context:
     opt_state: The state of the optimizer prior to the update. (available after
       the backward pass, e.g. for metrics). The old state is chosen to be
       consistent with parameters which are also pre-update.
+    metric_states: The states of the metrics (after the backward pass)
+    summary_states: The states of the summaries (after the backward pass)
   """
 
   # These are always available:
@@ -80,6 +82,9 @@ class Context:
   grads: Any = None
   updates: Any = None
   opt_state: Any = None
+  # Become available after the metrics computation
+  metric_states: Any = None
+  summary_states: Any = None
 
   replace = dataclasses.replace
 
@@ -100,3 +105,19 @@ def from_state_and_batch(
 
   def flatten(self) -> dict[str, Any]:
     return kontext.flatten_with_path(self)
+
+  def get_aux(
+      self,
+      *,
+      return_losses: bool = False,
+      return_metrics: bool = False,
+      return_summaries: bool = False,
+  ) -> train_step.Auxiliaries:
+    """Returns the auxiliaries for the step."""
+    from kauldron.train import train_step  # pylint: disable=g-import-not-at-top
+
+    return train_step.Auxiliaries(
+        loss_states=self.loss_states if return_losses else None,
+        metric_states=self.metric_states if return_metrics else None,
+        summary_states=self.summary_states if return_summaries else None,
+    )

kauldron/train/metric_writer.py

@@ -165,7 +165,6 @@ def write_step_metrics(
       *,
       step: int,
       aux: train_step.Auxiliaries,
-      model_with_aux: train_step.ModelWithAux,
       schedules: Mapping[str, optax.Schedule],
       log_summaries: bool,
       timer: Optional[chrono_utils.Chrono] = None,
@@ -201,15 +200,7 @@ def write_step_metrics(
 
     if log_summaries:
       with jax.transfer_guard("allow"):
-        # TODO(klausg): remove once all summaries are migrated to new protocol
-        # image summaries
-        image_summaries_old = {
-            name: summary.get_images(**aux.summary_kwargs[name])
-            for name, summary in model_with_aux.summaries.items()
-            if isinstance(summary, summaries.ImageSummary)
-        }
-
-        image_summaries = image_summaries_old | {
+        image_summaries = {
             name: value
             for name, value in aux_result.summary_values.items()
             if isinstance(value, Float["n h w #3"])
@@ -586,7 +577,6 @@ def write_step_metrics(
       *,
       step: int,
       aux: train_step.Auxiliaries,
-      model_with_aux: train_step.ModelWithAux,
       schedules: Mapping[str, optax.Schedule],
       log_summaries: bool,
       timer: Optional[chrono_utils.Chrono] = None,

kauldron/train/train_lib.py

@@ -136,7 +136,6 @@ def train_impl(
             step=i,
             aux=aux,
             schedules=trainer.schedules,
-            model_with_aux=trainstep.model_with_aux,
             timer=chrono,
             log_summaries=log_summaries,
         )