Remove ModelWithAux in trainstep

PiperOrigin-RevId: 699160009

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