Create objects for value functions for the estimator object and the corresponding implicit ggn estimator

PiperOrigin-RevId: 805859317

Author: KfacJaxDev

kfac_jax/_src/optimizer.py

@@ -545,12 +545,23 @@ def __init__(
         auto_register_tags=use_automatic_registration,
         auto_register_kwargs=auto_register_kwargs,
     )
+
+    self._value_func_for_estimator = (
+        self._value_func if value_func_for_estimator is None else
+        value_func_for_estimator)
+
     self._implicit = curvature_estimator.ImplicitExactCurvature(
         self._value_func,
         params_index=self._params_index,
         batch_size_extractor=batch_size_extractor,
     )
 
+    self._implicit_estimator = curvature_estimator.ImplicitExactCurvature(
+        self._value_func_for_estimator,
+        params_index=self._params_index,
+        batch_size_extractor=batch_size_extractor,
+    )
+
     # Each subclass should call finalize on its own, so this gets called only
     # for instances of exactly this class type.
     if type(self) == Optimizer:  # pylint: disable=unidiomatic-typecheck

kfac_jax/_src/tag_graph_matcher.py

@@ -1055,6 +1055,69 @@ def clean_jaxpr(
   return to_jaxpr_or_closed_jaxpr(closed_jaxpr, jaxpr)
 
 
+def clean_layer_tags_jaxpr(
+    jaxpr: J,
+    only_remove_auto_tags: bool = False,
+) -> tuple[J, tuple[tags.LayerTagEqn | JaxprEqn, ...]]:
+  """Removes layer tags from a Jaxpr."""
+
+  closed_jaxpr = to_closed_jaxpr(jaxpr)
+  eqns = []
+  layer_tag_eqns = []
+  var_map = {}
+
+  for eqn in closed_jaxpr.jaxpr.eqns:
+    if isinstance(eqn.primitive, tags.LayerTag) and (
+        not only_remove_auto_tags
+        or (
+            eqn.params["meta"].name is not None
+            and "Auto" in eqn.params["meta"].name
+        )
+    ):
+      for ind1, ind2 in enumerate(eqn.params["meta"].outputs_index):
+        var_map[eqn.outvars[ind1]] = eqn.invars[ind2]
+    else:
+      eqns.append(eqn)
+    if isinstance(eqn.primitive, tags.LayerTag):
+      layer_tag_eqns.append(eqn)
+
+  def remap_input_vars(
+      eqns: list[JaxprEqn], var_map: dict[jex.core.Var, jex.core.Var]
+  ) -> list[JaxprEqn]:
+    """Remaps the input variables of a JaxprEqn.
+
+    Args:
+      eqns: The list of JaxprEqns to remap.
+      var_map: A mapping from variables to new variables.
+
+    Returns:
+      A new list of JaxprEqns with remapped input variables.
+    """
+    eqns_new = []
+    for eqn in eqns:
+      new_invars = []
+      for var in eqn.invars:
+        if not isinstance(var, jex.core.Literal) and var in var_map.keys():
+          new_invars.append(var_map[var])
+        else:
+          new_invars.append(var)
+      eqns_new.append(eqn.replace(invars=new_invars))
+    return eqns_new
+
+  eqns_new = remap_input_vars(eqns, var_map)
+  layer_tag_eqns_new = remap_input_vars(layer_tag_eqns, var_map)
+
+  closed_jaxpr = ClosedJaxpr(
+      jaxpr=closed_jaxpr.jaxpr.replace(eqns=eqns_new),
+      consts=closed_jaxpr.consts,
+  )
+
+  return (
+      to_jaxpr_or_closed_jaxpr(closed_jaxpr, jaxpr),
+      tuple(layer_tag_eqns_new),
+  )
+
+
 # Prototype for clean_jaxpr using JAX's dce_jaxpr. Doesn't work because
 # dce_jaxpr will remove any equations with no used outputs, regardless of the
 # dce_rule for that equation's primitive. Adding an "effect" to loss/layer

kfac_jax/_src/tracer.py

@@ -225,7 +225,11 @@ def __init__(
     self.consts = consts
     self.in_tree = in_tree
     self.params_index = params_index
-    self.layer_tags, self.loss_tags = extract_tags(jaxpr)
+    # clean jaxpr of layer tags at this point
+    closed_jaxpr = jex.core.ClosedJaxpr(jaxpr=self.jaxpr, consts=self.consts)
+    self.jaxpr, self.layer_tags = tgm.clean_layer_tags_jaxpr(closed_jaxpr)
+    self.jaxpr = self.jaxpr.jaxpr
+    _, self.loss_tags = extract_tags(self.jaxpr)
     name_layer_tags(self.layer_tags)
     self.layer_tags, self.layer_indices = order_layer_tags(
         params_vars_flat=self.params_vars_flat,
@@ -579,7 +583,7 @@ def _loss_tags_vjp(
     p_jaxpr: ProcessedJaxpr,
     primal_func_args: FuncArgs,
 ) -> LossTagsVjp:
-  """Computes a (backward-mode) vector-Jacobian product w.r.t. all loss tags.
+  """Computes a (backward-mode) vector-Jacobian product for the vector of losses given by the loss tags.
 
   The function has similar interface to :func:`jax.vjp`. It takes as inputs the
   concrete values of the primals at which the Jacobian will be evaluated. It
@@ -651,7 +655,7 @@ def _loss_tags_jvp(
     primal_func_args: FuncArgs,
     params_tangents: Params,
 ) -> LossTagsJvp:
-  """Computes a (forward-mode) Jacobian-vector product w.r.t. all loss tags.
+  """Computes a (forward-mode) Jacobian-vector product for the losses given by the loss tags.
 
   The function has similar interface to :func:`jax.jvp`. It takes as inputs the
   concrete values of the primals at which the Jacobian will be evaluated at and
@@ -795,7 +799,7 @@ def _layer_tag_vjp(
   Args:
     processed_jaxpr: The :class:`~ProcessedJaxpr` representing the model
       function. This must include at least one loss tag.
-    primal_func_args: The primals at which to evaluate the Hessian.
+    primal_func_args: The primals at which to evaluate the Jacobian.
 
   Returns:
     The computed ``losses`` and ``vjp_func`` pair.
@@ -824,14 +828,12 @@ def forward() -> tuple[Array, ...]:
     # Loop through equations and evaluate them
     num_losses_passed = 0
     for eqn in processed_jaxpr.jaxpr.eqns:
-
-      write(eqn.outvars, tgm.eval_jaxpr_eqn(eqn, read(eqn.invars)))
-
       if isinstance(eqn.primitive, tags.LossTag):
         num_losses_passed += 1
         if num_losses_passed == len(processed_jaxpr.loss_tags):
           break
-
+      else:
+        write(eqn.outvars, tgm.eval_jaxpr_eqn(eqn, read(eqn.invars)))
     assert num_losses_passed == len(processed_jaxpr.loss_tags)
 
     return tuple(read(layer_input_vars))
@@ -884,7 +886,6 @@ def write(variables: list[jex.core.Var], values: list[Array]) -> None:
     for eqn in processed_jaxpr.jaxpr.eqns:
 
       input_values = read(eqn.invars)
-      write(eqn.outvars, tgm.eval_jaxpr_eqn(eqn, input_values))
 
       if isinstance(eqn.primitive, tags.LossTag):
         loss: LossFunction = tags.loss_eqn_construct_loss(eqn, *input_values)
@@ -896,6 +897,8 @@ def write(variables: list[jex.core.Var], values: list[Array]) -> None:
 
         if num_losses_passed == len(processed_jaxpr.loss_tags):
           break
+      else:
+        write(eqn.outvars, tgm.eval_jaxpr_eqn(eqn, input_values))
 
     assert num_losses_passed == len(processed_jaxpr.loss_tags)