Use public aliases to define MultiTrainStep

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,8 @@
 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 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
 

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