Add function to remove layer tags from jaxpr

PiperOrigin-RevId: 804950216

Author: KfacJaxDev

kfac_jax/_src/tag_graph_matcher.py

@@ -1055,6 +1055,62 @@ 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,
+) -> J:
+  """Removes layer tags from a Jaxpr."""
+
+  closed_jaxpr = to_closed_jaxpr(jaxpr)
+  eqns = []
+  var_map = {}
+
+  for eqn in closed_jaxpr.jaxpr.eqns:
+    if isinstance(eqn.primitive, tags.LayerTag) and (
+        not only_remove_auto_tags or "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)
+
+  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))
+
+  closed_jaxpr = ClosedJaxpr(
+      jaxpr=closed_jaxpr.jaxpr.replace(eqns=eqns_new),
+      consts=closed_jaxpr.consts,
+  )
+
+  return to_jaxpr_or_closed_jaxpr(closed_jaxpr, jaxpr)
+
+
+def clean_layer_tags_func(
+    func: utils.Func,
+    func_args: utils.FuncArgs,
+    only_remove_auto_tags: bool = False,
+) -> utils.Func:
+  """Removes layer tags from a function."""
+  typed_jaxpr = jax.make_jaxpr(func)(*func_args)
+  jaxpr_clean = clean_layer_tags_jaxpr(typed_jaxpr, only_remove_auto_tags)
+
+  def func_clean(*args):
+    # eval_jaxpr takes the jaxpr, its constants (literals), and arguments
+    flattened_args = jax.tree_util.tree_leaves(args)
+    return jax.core.eval_jaxpr(
+        jaxpr_clean.jaxpr, jaxpr_clean.literals, *flattened_args
+    )
+
+  return func_clean
+
+
 # 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

@@ -321,6 +321,11 @@ def make_from_func(
     func_args = tuple(func_args)
 
     if auto_register_tags:
+      # Always remove existing auto layer tags before re-running the automatic
+      # registration.
+      func = tgm.clean_layer_tags_func(
+          func, func_args, only_remove_auto_tags=True
+      )
       func = tgm.auto_register_tags(
           func=func,
           func_args=func_args,
@@ -824,14 +829,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 +887,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 +898,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)