Add the option to process params with a custom function before writin…

…g them.

PiperOrigin-RevId: 662889423

kauldron/train/train_lib.py

@@ -74,7 +74,7 @@ def train_impl(
     # `NoopWriter`, the second one is mandatory in order to re-load the model
     # for evals/inference.
     writer.write_config(trainer.raw_cfg)
-    writer.write_param_overview(initial_step, state.params)
+    writer.write_param_overview(initial_step, trainer.process_params_for_writing(state.params))
     writer.write_element_spec(initial_step, ds_iter.element_spec)
     writer.write_context_structure(initial_step, trainer)
 

kauldron/train/trainer_lib.py

@@ -20,7 +20,7 @@
 import dataclasses
 import functools
 import typing
-from typing import Any, Optional, Self
+from typing import Any, Callable, Optional, Self
 
 from etils import edc
 from etils import epath
@@ -237,6 +237,10 @@ class Trainer(config_util.BaseConfig):
       default=None, repr=False
   )
 
+  process_params_for_writing: Callable[[Any], Any] = dataclasses.field(
+      default=lambda x: x, repr=False
+  )
+
   def __post_init__(self):
     # It's convenient to set `cfg.evals = None` to disable evaluation
     for name, default_factory in {