Merge pull request #30062 from jax-ml:auto-dce-dynamic-jaxpr-trace

PiperOrigin-RevId: 781744820

Author: Google-ML-Automation

jax/_src/ad_checkpoint.py

@@ -448,11 +448,14 @@ def f_(*args):
     return f(*args, **kwargs)
 
   debug_info = api_util.debug_info("saved_residuals", f, args, kwargs)
-  out = api.make_jaxpr(lambda *args: api.linearize(f_, *args)[1],
+  out = api.make_jaxpr(lambda *args: api.linearize(f_, *args),
                        return_shape=True)(*in_leaves)
   assert isinstance(out, tuple)
-  jaxpr_, out_shape = out
+  jaxpr_, out_shape_ = out
   jaxpr = jaxpr_.jaxpr
+  out_shape = out_shape_[1]
+  num_res = tree_structure(out_shape).num_leaves
+  jaxpr = jaxpr.replace(outvars=jaxpr.outvars[len(jaxpr.outvars) - num_res:])
   out_tree = lambda: tree_structure(out_shape)
   assert len(jaxpr.invars) == len(in_leaves)
   return _saved_residuals(jaxpr, debug_info.arg_names)

jax/_src/checkify.py

@@ -472,6 +472,7 @@ def _reduce_any_error(error: Error):
 ## check_p primitive
 
 check_p = core.Primitive('check')
+check_p.is_effectful = lambda _: True  # type: ignore
 check_p.multiple_results = True  # zero results
 
 
@@ -841,12 +842,12 @@ def new_body_f(*c_consts_and_vals):
     c_consts, vals = split_list(c_consts_and_vals, [c_consts_num])
     out = body_f(*vals)
     # This checks if the next cond application will error
-    _ = cond_f(*c_consts, *out)
+    lax.dce_sink(cond_f(*c_consts, *out))
     return out
   new_body_f_ = lu.wrap_init(new_body_f, debug_info=body_jaxpr.jaxpr.debug_info)
   c_consts_avals = cond_jaxpr.in_avals[:c_consts_num]
-  jaxpr, _, () = pe.trace_to_jaxpr_dynamic(new_body_f_, [*c_consts_avals,
-                                                             *body_jaxpr.in_avals])
+  jaxpr, _, () = pe.trace_to_jaxpr_dynamic(
+      new_body_f_, [*c_consts_avals, *body_jaxpr.in_avals])
   closed_jaxpr = pe.close_jaxpr(jaxpr)
   err_vals, err_tree = jtu.tree_flatten(error)
   err_vals = map(core.get_aval, err_vals)
@@ -1232,9 +1233,8 @@ def checked_fun(*args, **kwargs):
     closed_f = lambda: f(*args, **kwargs)
     # stage:
     debug = api_util.debug_info("checkify", f, args, kwargs)
-    fun_, out_tree = api_util.flatten_fun(lu.wrap_init(closed_f,
-                                                       debug_info=debug),
-                                          in_tree)
+    fun_, out_tree = api_util.flatten_fun(
+        lu.wrap_init(closed_f, debug_info=debug), in_tree)
     jaxpr_, _, consts = pe.trace_to_jaxpr_dynamic(fun_, ())
     jaxpr = pe.close_jaxpr(pe.convert_constvars_jaxpr(jaxpr_))
     # checkify:

jax/_src/core.py

@@ -605,6 +605,10 @@ def def_effectful_abstract_eval(self, effectful_abstract_eval):
     self.abstract_eval = effectful_abstract_eval
     return effectful_abstract_eval
 
+  def def_effectful_abstract_eval2(self, abstract_eval):
+    self.abstract_eval = _generic_effectful_abstract_eval(abstract_eval)
+    return abstract_eval
+
   def def_bind_with_trace(self, bind_with_trace):
     self.bind_with_trace = bind_with_trace
     return bind_with_trace
@@ -629,6 +633,19 @@ def abstract_eval_(*args, **kwargs):
     return abstract_eval(*args, **kwargs), no_effects
   return abstract_eval_
 
+class GenericEffect(Effect):
+  pass
+generic_effect = GenericEffect()
+generic_effect_set = {generic_effect}
+effects.lowerable_effects.add_type(GenericEffect)
+effects.control_flow_allowed_effects.add_type(GenericEffect)
+effects.custom_derivatives_allowed_effects.add_type(GenericEffect)
+
+def _generic_effectful_abstract_eval(abstract_eval):
+  def abstract_eval_(*args, **kwargs):
+    return abstract_eval(*args, **kwargs), generic_effect_set
+  return abstract_eval_
+
 # -------------------- lifting --------------------
 
 # TODO(mattjj): replace this approach with a primitive-keyed table of rules

jax/_src/interpreters/ad.py

@@ -93,7 +93,7 @@ def linearize_subtrace(_f: Callable, _store: lu.Store, _tag: core.TraceTag,
                        *primals, **params):
   source_info = source_info_util.current()
   with core.take_current_trace() as parent_trace:
-    tangent_trace = pe.DynamicJaxprTrace(debug_info)
+    tangent_trace = pe.DynamicJaxprTrace(debug_info, auto_dce=True)
     tangent_trace.tag = _tag
     linearize_trace = LinearizeTrace(parent_trace, tangent_trace, tag=_tag)
     tracers = [LinearizeTracer(linearize_trace, p,
@@ -167,7 +167,7 @@ def _linearize_jaxpr(
 ) -> tuple[core.ClosedJaxpr, int, Sequence[bool], Sequence[int | None], core.ClosedJaxpr]:
   dbg = jaxpr.jaxpr.debug_info
   primal_trace = pe.DynamicJaxprTrace(dbg)
-  tangent_trace = pe.DynamicJaxprTrace(dbg)
+  tangent_trace = pe.DynamicJaxprTrace(dbg, auto_dce=True)
   lin_trace = LinearizeTrace(primal_trace, tangent_trace)
   tangent_trace.tag = lin_trace.tag
 
@@ -227,7 +227,7 @@ def direct_linearize(traceable: lu.WrappedFun,
                      primals, kwargs, *, has_aux=False, tag=None):
   with core.take_current_trace() as parent_trace:
     source_info = source_info_util.current()
-    tangent_trace = pe.DynamicJaxprTrace(traceable.debug_info)
+    tangent_trace = pe.DynamicJaxprTrace(traceable.debug_info, auto_dce=True)
     tangents = [tangent_trace.new_arg(get_aval(p).to_tangent_aval(), source_info) for p in primals]
     tangents = [Zero.from_primal_value(t) if dtype(t) == float0 else t for t in tangents]
     linearize_trace = LinearizeTrace(parent_trace, tangent_trace, tag=tag)