Internal Change

PiperOrigin-RevId: 805859317

Author: KfacJaxDev

examples/optimizers.py

@@ -114,6 +114,7 @@ def create_optimizer(
     total_steps: int | None,
     total_epochs: float | None,
     schedule_free_config: config_dict.ConfigDict,
+    model_func_for_power_iteration: kfac_jax.optimizer.ValueFunc | None = None,
 ) -> optax_wrapper.OptaxWrapper | kfac_jax.Optimizer:
   """Creates an optimizer from the provided configuration."""
 
@@ -154,6 +155,11 @@ def create_optimizer(
               **kwargs[sched_name]
               )
 
+      if model_func_for_power_iteration is not None:
+        kwargs["value_func_for_power_iteration"] = (
+            model_func_for_power_iteration
+        )
+
     return kfac_jax.Optimizer(
         value_and_grad_func=value_and_grad_func,
         l2_reg=l2_reg,

examples/training.py

@@ -149,6 +149,8 @@ class SupervisedExperiment(abc.ABC):
       computation of the loss of the model for the estimator.
     eval_splits: Evaluation splits of the evaluation dataset loader.
     batch_size: An instance of `ExperimentBatchSizes`.
+    model_func_for_power_iteration: A function that allows a different
+      computation of the loss of the model for power iteration.
     init_parameters_func: A function that initializes the parameters and
       optionally the state of the model if it has one.
     params_init: A function that initializes the model parameters.
@@ -173,6 +175,9 @@ def __init__(
       model_func_for_estimator: kfac_jax.optimizer.ValueFunc | None = None,
       eval_splits: tuple[str, ...] = ("train", "test"),
       batch_size_calculator_ctor: BatchSizeCalculatorCtor = BatchSizeCalculator,
+      model_func_for_power_iteration: (
+          kfac_jax.optimizer.ValueFunc | None
+      ) = None,
   ):
     """Initializes experiment.
 
@@ -193,6 +198,8 @@ def __init__(
       eval_splits: Evaluation splits of the evaluation dataset loader.
       batch_size_calculator_ctor: A constructor function to create a batch size
         calculator.
+      model_func_for_power_iteration: A function that allows a different
+        computation of the loss of the model for power iteration.
     """
     self.mode = mode
     self.init_rng, seed_rng = jax.random.split(init_rng)
@@ -215,6 +222,7 @@ def __init__(
     self.params_init = jax.pmap(init_parameters_func, axis_name="kfac_axis")
     self.model_loss_func = model_loss_func
     self.model_func_for_estimator = model_func_for_estimator
+    self.model_func_for_power_iteration = model_func_for_power_iteration
 
     self.train_model_func = functools.partial(
         self.model_loss_func, is_training=True
@@ -467,6 +475,9 @@ def create_optimizer(
         total_steps=self.config.training.steps,
         total_epochs=self.config.training.epochs,
         schedule_free_config=self._schedule_free_config,
+        model_func_for_power_iteration=functools.partial(
+            self.model_func_for_power_iteration, is_training=True
+        ) if self.model_func_for_power_iteration is not None else None,
     )
 
   def maybe_initialize_state(self):

kfac_jax/_src/optimizer.py

@@ -545,6 +545,7 @@ def __init__(
         auto_register_tags=use_automatic_registration,
         auto_register_kwargs=auto_register_kwargs,
     )
+
     self._implicit = curvature_estimator.ImplicitExactCurvature(
         self._value_func,
         params_index=self._params_index,