[no-thunks] Reduce stack frames in cond and friends

Also deprecate the very old form of `cond`:
`cond(predicate, true_arg, true_fun, false_arg, false_fun)`.

Author: dougalm

jax/_src/interpreters/partial_eval.py

@@ -2400,7 +2400,7 @@ def trace_to_jaxpr(
     in_tree: PyTreeDef,
     in_avals_flat: Sequence[AbstractValue | core.AvalQDD],
     debug_info: core.DebugInfo
-) -> tuple[Jaxpr, PyTreeDef, list[Any]]:
+) -> tuple[ClosedJaxpr, PyTreeDef, list[Any]]:
   config.enable_checks.value and debug_info.assert_arg_names(len(in_avals_flat))
   parent_trace = core.trace_ctx.trace
   trace = DynamicJaxprTrace(debug_info, parent_trace=parent_trace)
@@ -2424,6 +2424,8 @@ def trace_to_jaxpr(
     del trace, fun, in_tracers_flat, in_tracers, out_tracers, ans, ans_flat
 
   config.enable_checks.value and core.check_jaxpr(jaxpr)
+  # TODO(dougalm): remove this once we merge Jaxpr and ClosedJaxpr
+  jaxpr = close_jaxpr(convert_constvars_jaxpr(jaxpr))
   return jaxpr, out_tree, consts
 
 

jax/_src/lax/control_flow/__init__.py

@@ -50,9 +50,6 @@
 # Private utilities used elsewhere in JAX
 # TODO(sharadmv): lift them into a more common place
 from jax._src.lax.control_flow.common import (
-    _initial_style_open_jaxpr as _initial_style_open_jaxpr,
-    _initial_style_jaxpr as _initial_style_jaxpr,
-    _initial_style_jaxprs_with_common_consts as _initial_style_jaxprs_with_common_consts,
     _check_tree_and_avals as _check_tree_and_avals,
 
 )

jax/_src/lax/control_flow/common.py

@@ -43,78 +43,54 @@ def _typecheck_param(prim, param, name, msg_required, pred):
     msg = sep.join([msg, param_str])
     raise core.JaxprTypeError(msg)
 
-# TODO(dougalm): this is a silly wrapper now. Delete it.
-@weakref_lru_cache
-def _initial_style_open_jaxpr(fun: Callable,
-                              in_tree: PyTreeDef,
-                              in_avals: Sequence[core.AbstractValue | core.AvalQDD],
-                              debug_info: core.DebugInfo):
-  jaxpr, out_tree, consts = pe.trace_to_jaxpr(fun, in_tree, in_avals, debug_info)
-  return jaxpr, consts, out_tree
-
-# TODO(dougalm): Delete. Make `trace_to_jaxpr` do the jaxpr-closing thing instead.
-@weakref_lru_cache
-def _initial_style_jaxpr(fun: Callable,
-                         in_tree: PyTreeDef,
-                         in_avals: Sequence[core.AbstractValue],
-                         debug_info: core.DebugInfo) -> tuple[core.ClosedJaxpr, Sequence[Any], PyTreeDef]:
-  jaxpr, consts, out_tree = _initial_style_open_jaxpr(
-      fun, in_tree, in_avals, debug_info)
-  closed_jaxpr = pe.close_jaxpr(pe.convert_constvars_jaxpr(jaxpr))
-  return closed_jaxpr, consts, out_tree
-
-def _initial_style_jaxprs_with_common_consts(
-    funs: Sequence[Callable],
-    in_tree: PyTreeDef, in_avals: Sequence[core.AbstractValue | core.AvalQDD],
-    debug_infos: Sequence[core.DebugInfo]):
-  jaxpr_data = [_initial_style_open_jaxpr(fn, in_tree, in_avals, debug_info)
-                for fn, debug_info in zip(funs, debug_infos)]
-  if not jaxpr_data: return [], [], []
-  jaxprs, all_consts, all_out_trees = zip(*jaxpr_data)
-
+# TODO(dougalm): this seems way too complicated. Why not allow different consts for each
+# branch of a switch?
+def _merge_common_consts(
+    jaxprs: Sequence[core.Jaxpr],
+    all_consts: Sequence[Sequence[Any]]
+    ) -> tuple[Sequence[core.ClosedJaxpr], Sequence[Any]]:
   # Jaxprs must share consts, so we concat consts and pad the jaxprs' constvars.
   lens = map(len, all_consts)
   consts = [c for cs in all_consts for c in cs]
   avalqdds = tuple(map(core.cur_aval_qdd, consts))
-  jaxprs = [_pad_constvars(jaxpr, avalqdds[:sum(lens[:i])], avalqdds[sum(lens[:i+1]):])
-            for i, jaxpr in enumerate(jaxprs)]
+  num_constss = [len(cs) for cs in all_consts]
+  jaxprs = [_pad_constvars(jaxpr, num_consts, avalqdds[:sum(lens[:i])], avalqdds[sum(lens[:i+1]):])
+            for i, (jaxpr, num_consts) in enumerate(zip(jaxprs, num_constss))]
   # De-duplicate shared constants.
   const_ids = tuple(id(c) for c in consts)
   seen = set()
-  consts = [c for c in consts if id(c) not in seen and not seen.add(id(c))]  # type: ignore
-  jaxprs = [_dedup_consts(jaxpr, const_ids) for jaxpr in jaxprs]
-
-  closed_jaxprs = [pe.close_jaxpr(pe.convert_constvars_jaxpr(jaxpr))
-                   for jaxpr in jaxprs]
-  return closed_jaxprs, consts, all_out_trees
+  dd_consts = [c for c in consts if id(c) not in seen and not seen.add(id(c))]  # type: ignore
+  jaxprs = [_dedup_consts(jaxpr, len(consts), const_ids) for jaxpr in jaxprs]
+  return jaxprs, dd_consts
 
 @weakref_lru_cache
-def _pad_constvars(jaxpr: core.Jaxpr, left: tuple[core.AvalQDD, ...],
-                   right: tuple[core.AbstractValue, ...]) -> core.Jaxpr:
+def _pad_constvars(jaxpr: core.ClosedJaxpr, num_consts: int,
+                   left: tuple[core.AvalQDD, ...],
+                   right: tuple[core.AbstractValue, ...]) -> core.ClosedJaxpr:
   def make_var(aq):
     return core.Var(aq.aval, initial_qdd=aq.qdd, final_qdd=aq.qdd)
-  constvars = [*map(make_var, left), *jaxpr.constvars, *map(make_var, right)]
-  effs = pe._renumber_effects([*constvars, *jaxpr.invars],
-                              [*jaxpr.constvars, *jaxpr.invars], jaxpr.effects)
-  jaxpr = jaxpr.replace(constvars=constvars, effects=effs)
+  invars = [*map(make_var, left), *jaxpr.invars[:num_consts],
+            *map(make_var, right), *jaxpr.invars[num_consts:]]
+  effs = pe._renumber_effects(invars, jaxpr.invars, jaxpr.effects)
+  jaxpr = jaxpr.replace(jaxpr=jaxpr.jaxpr.replace(invars=invars, effects=effs))
   config.enable_checks.value and core.check_jaxpr(jaxpr)
   return jaxpr
 
 @weakref_lru_cache
-def _dedup_consts(jaxpr, const_ids):
+def _dedup_consts(jaxpr, num_consts, const_ids):
   newvars = {}
   canonicalize = {v: newvars.setdefault(constid, v)
-                  for constid, v in zip(const_ids, jaxpr.constvars)}
+                  for constid, v in zip(const_ids, jaxpr.invars[:num_consts])}
   eqns = [e.replace(invars=[canonicalize.get(x, x) if isinstance(x, core.Var)
                             else x for x in e.invars]) for e in jaxpr.eqns]
   outvars = [canonicalize.get(x, x) if isinstance(x, core.Var) else x
              for x in jaxpr.outvars]
-  constvars = list(newvars.values())
-  effs = pe._renumber_effects(
-      [*constvars, *jaxpr.invars],
-      [*map(canonicalize.get, jaxpr.constvars), *jaxpr.invars], jaxpr.effects)
-  jaxpr = jaxpr.replace(constvars=constvars, eqns=eqns, outvars=outvars,
-                        effects=effs)
+  invars = [*list(newvars.values()), *jaxpr.invars[num_consts:]]
+  effs = pe._renumber_effects(invars,
+      [*map(canonicalize.get, jaxpr.invars[:num_consts]), *jaxpr.invars[num_consts:]],
+      jaxpr.effects)
+  jaxpr = jaxpr.replace(jaxpr=jaxpr.jaxpr.replace(invars=invars, eqns=eqns, outvars=outvars,
+                        effects=effs))
   config.enable_checks.value and core.check_jaxpr(jaxpr)
   return jaxpr
 

jax/_src/lax/control_flow/conditionals.py

@@ -53,7 +53,7 @@
 import numpy as np
 
 from jax._src.lax.control_flow.common import (
-    _avals_short, _typecheck_param, _initial_style_jaxprs_with_common_consts,
+    _avals_short, _typecheck_param, _merge_common_consts,
     _make_closed_jaxpr, _prune_zeros)
 
 map, unsafe_map = safe_map, map
@@ -149,8 +149,10 @@ def _switch_internal(
   if config.mutable_array_checks.value:
     api_util.check_no_aliased_ref_args(lambda: dbgs[0], ops_avals, ops)
 
-  jaxprs, consts, out_trees = _initial_style_jaxprs_with_common_consts(
-      branches, ops_tree, ops_avals, dbgs)
+  jaxprs_, out_trees, all_consts = zip(*[pe.trace_to_jaxpr(
+      branch, ops_tree, ops_avals, dbg) for branch, dbg in zip(branches, dbgs)])
+  jaxprs, consts = _merge_common_consts(jaxprs_, all_consts)
+
   if config.mutable_array_checks.value:
     api_util._check_no_aliased_closed_over_refs(dbgs[0], (*jaxprs[0].consts, *consts), ops)
   for i, (out_tree, jaxpr) in enumerate(zip(out_trees[1:], jaxprs[1:])):
@@ -184,7 +186,7 @@ def _switch_internal(
   return tree_unflatten(out_trees[0], out)
 
 @partial(api_boundary, repro_api_name="jax_cond")
-def _cond(pred, true_fun: Callable, false_fun: Callable, *operands,
+def cond(pred, true_fun: Callable, false_fun: Callable, *operands,
           operand=_no_operand_sentinel):
   """Conditionally apply ``true_fun`` or ``false_fun``.
 
@@ -270,14 +272,16 @@ def cond(pred, true_fun, false_fun, *operands):
   if config.mutable_array_checks.value:
     api_util.check_no_aliased_ref_args(lambda: dbg_true_fun, ops_avals, ops)
   dbg_false_fun = api_util.debug_info("cond", false_fun, operands, {})
-  jaxprs, consts, out_trees = _initial_style_jaxprs_with_common_consts(
-      (true_fun, false_fun), ops_tree, ops_avals,
-      [dbg_true_fun, dbg_false_fun])
-  true_jaxpr, false_jaxpr = jaxprs
+
+  true_jaxpr_, out_tree, true_consts  = pe.trace_to_jaxpr(
+      true_fun, ops_tree, ops_avals, dbg_true_fun)
+  false_jaxpr_, false_out_tree, false_consts = pe.trace_to_jaxpr(
+      false_fun, ops_tree, ops_avals, dbg_false_fun)
+  (true_jaxpr, false_jaxpr), consts = _merge_common_consts(
+      (true_jaxpr_, false_jaxpr_), (true_consts, false_consts))
   if config.mutable_array_checks.value:
     api_util._check_no_aliased_closed_over_refs(dbg_true_fun, (*true_jaxpr.consts, *consts), ops)
 
-  out_tree, false_out_tree = out_trees
   if any(isinstance(out_aval, AbstractRef) for out_aval in
          true_jaxpr.out_avals + false_jaxpr.out_avals):
     raise ValueError("Cannot return `Ref`s from `cond`.")
@@ -399,48 +403,6 @@ def _capitalize(s):
   # s.capitalize() converts s[1:] to lowercase which we don't want.
   return s[0].capitalize() + s[1:]
 
-@api_boundary
-@functools.wraps(_cond)
-def cond(*args, **kwargs):
-  # detect an attempt to call the former, deprecated cond
-  try:
-    ba = inspect.signature(_cond_with_per_branch_args).bind(*args, **kwargs)
-  except TypeError:
-    pass
-  else:
-    assert not ba.kwargs  # no catch-all **kwargs in _cond_with_per_branch
-    _, true_operand, true_fun, false_operand, false_fun = ba.args
-    if callable(true_operand) and callable(true_fun):
-      # treat this as modern cond (with two operands)
-      return _cond(*args, **kwargs)
-    if callable(true_fun) and callable(false_fun):
-      return _cond_with_per_branch_args(*ba.args)
-
-  return _cond(*args, **kwargs)
-
-@partial(api_boundary, repro_api_name="jax_cond_with_per_branch_args")
-def _cond_with_per_branch_args(pred,
-                               true_operand, true_fun: Callable,
-                               false_operand, false_fun: Callable):
-  """Conditionally apply ``true_fun`` or ``false_fun``.
-
-  Has equivalent semantics to this Python implementation::
-
-    def cond(pred, true_operand, true_fun, false_operand, false_fun):
-      if pred:
-        return true_fun(true_operand)
-      else:
-        return false_fun(false_operand)
-
-  Pred has to be a scalar type, collection types (list, tuple) are not supported
-  """
-  if not (callable(true_fun) and callable(false_fun)):
-    raise TypeError("lax.cond: true_fun and false_fun arguments should be callable.")
-  return _cond(pred,
-               lambda op: true_fun(op[0]),
-               lambda op: false_fun(op[1]),
-               (true_operand, false_operand))
-
 def _join_cond_effects(branches: Sequence[core.ClosedJaxpr]) -> effects.Effects:
   joined_effects = set()
   for b in branches:

jax/_src/lax/control_flow/loops.py

@@ -50,7 +50,7 @@
 from jax._src.lax import slicing
 from jax._src.lax import windowed_reductions
 from jax._src.lax.control_flow.common import (
-    _avals_short, _initial_style_jaxpr, _prune_zeros, _typecheck_param,
+    _avals_short, _prune_zeros, _typecheck_param,
     _make_closed_jaxpr)
 from jax._src.lax.other import logaddexp
 from jax._src.pjit import auto_axes, PartitionSpec as P, reshard
@@ -281,9 +281,8 @@ def _create_jaxpr(init):
     init_flat, init_tree = tree_flatten(init)
     in_flat, in_tree = tree_flatten((init, xs))
     carry_avals = tuple(_map(core.get_aval, init_flat))
-    open_jaxpr, out_tree, consts = pe.trace_to_jaxpr(
+    jaxpr, out_tree, consts = pe.trace_to_jaxpr(
         f, in_tree, (*carry_avals, *x_avals), debug_info=dbg_body)
-    jaxpr = pe.close_jaxpr(pe.convert_constvars_jaxpr(open_jaxpr))
     if config.mutable_array_checks.value:
       _check_no_aliased_closed_over_refs(dbg_body, (*jaxpr.consts, *consts), in_flat)
     out_tree_children = out_tree.children()
@@ -1712,10 +1711,10 @@ def _create_jaxpr(init_val):
     init_vals, in_tree = tree_flatten((init_val,))
     init_avals = tuple(_map(core.get_aval, init_vals))
     cond_dbg = api_util.debug_info("while_cond", cond_fun, (init_val,), {})
-    cond_jaxpr, cond_consts, cond_tree = _initial_style_jaxpr(
+    cond_jaxpr, cond_tree, cond_consts = pe.trace_to_jaxpr(
         cond_fun, in_tree, init_avals, cond_dbg)
     body_dbg = api_util.debug_info("while_body", body_fun, (init_val,), {})
-    body_jaxpr, body_consts, body_tree = _initial_style_jaxpr(
+    body_jaxpr, body_tree, body_consts = pe.trace_to_jaxpr(
         body_fun, in_tree, init_avals, body_dbg)
     if not treedef_is_leaf(cond_tree) or len(cond_jaxpr.out_avals) != 1:
       msg = "cond_fun must return a boolean scalar, but got pytree {}."

jax/_src/lax/control_flow/solves.py

@@ -27,6 +27,7 @@
 from jax._src.interpreters import ad
 from jax._src.interpreters import batching
 from jax._src.interpreters import mlir
+from jax._src.interpreters import partial_eval as pe
 from jax._src.interpreters import pxla
 from jax._src.traceback_util import api_boundary
 from jax._src.tree_util import (tree_flatten, treedef_children, tree_leaves,
@@ -36,7 +37,6 @@
 
 from jax._src.lax.control_flow.common import (
     _check_tree,
-    _initial_style_jaxpr,
     )
 
 _map = safe_map
@@ -95,7 +95,7 @@ def custom_root(f: Callable,
   guess_flat, in_args_tree = tree_flatten((initial_guess,))
   guess_avals = tuple(_map(core.get_aval, guess_flat))
   f_debug = api_util.debug_info("custom_root", f, (initial_guess,), {})
-  f_jaxpr, f_consts, out_tree = _initial_style_jaxpr(
+  f_jaxpr, out_tree, f_consts = pe.trace_to_jaxpr(
       f, in_args_tree, guess_avals, f_debug)
 
   in_tree, = treedef_children(in_args_tree)
@@ -104,7 +104,7 @@ def custom_root(f: Callable,
   solve_debug = api_util.debug_info("custom_root solve", solve,
                                     (f, initial_guess), {},
                                     static_argnums=(0,))
-  solve_jaxpr, solve_consts, solution_tree = _initial_style_jaxpr(
+  solve_jaxpr, solution_tree, solve_consts = pe.trace_to_jaxpr(
       partial(solve, f), in_args_tree, guess_avals, solve_debug)
   _check_tree("solve", "initial_guess", solution_tree, in_tree, has_aux)
 
@@ -114,7 +114,7 @@ def linearize_and_solve(x, b):
 
   linearize_and_solve_dbg = api_util.debug_info("custom_root tangent_solve",
       tangent_solve, (initial_guess, initial_guess), {})
-  l_and_s_jaxpr, l_and_s_consts, out_tree = _initial_style_jaxpr(
+  l_and_s_jaxpr, out_tree, l_and_s_consts = pe.trace_to_jaxpr(
       linearize_and_solve, treedef_tuple((in_tree,) * 2), guess_avals * 2,
       linearize_and_solve_dbg)
   _check_tree("tangent_solve", "x", out_tree, in_tree, False)
@@ -268,15 +268,15 @@ def f_aux(x):
   matvec_debug = api_util.debug_info("custom_linear_solve",
                                      matvec, (b,), {})
   # no auxiliary data assumed for matvec
-  matvec_jaxpr, matvec_consts, out_tree = _initial_style_jaxpr(
+  matvec_jaxpr, out_tree, matvec_consts = pe.trace_to_jaxpr(
       _shape_checked(matvec, "matvec", False), in_args_tree, b_avals,
       matvec_debug)
   _check_tree("matvec", "b", out_tree, tree, False)
 
   solve_debug = api_util.debug_info("custom_linear_solve solve",
                                     solve, (matvec, b), {},
                                     static_argnums=(0,))
-  solve_jaxpr, solve_consts, out_tree = _initial_style_jaxpr(
+  solve_jaxpr, out_tree, solve_consts = pe.trace_to_jaxpr(
       _shape_checked(partial(solve, matvec), "solve", has_aux), in_args_tree, b_avals,
       solve_debug)
   _check_tree("solve", "b", out_tree, tree, has_aux)
@@ -294,11 +294,11 @@ def f_aux(x):
       vecmat_consts = matvec_consts
     else:
       vecmat = _transpose_one_output(matvec, b)
-      vecmat_jaxpr, vecmat_consts, out_tree = _initial_style_jaxpr(
+      vecmat_jaxpr, out_tree, vecmat_consts = pe.trace_to_jaxpr(
           vecmat, in_args_tree, b_avals, transpose_solve_debug)
       assert out_tree == tree
 
-    tr_solve_jaxpr, tr_solve_consts, out_tree = _initial_style_jaxpr(
+    tr_solve_jaxpr, out_tree, tr_solve_consts = pe.trace_to_jaxpr(
         _shape_checked(partial(transpose_solve, vecmat), "transpose_solve", has_aux),
         in_args_tree, b_avals, transpose_solve_debug)
     _check_tree("transpose_solve", "b", out_tree, tree, has_aux)

tests/api_test.py

@@ -7207,10 +7207,10 @@ def fun(x):
   def test_cond(self):
     def f(x):
       return lax.cond(x >= 0.,
+                      lambda xt, _: xt + x,
+                      lambda _, xf: xf - x,
                       x + 1.,
-                      lambda xt: xt + x,
-                      x + 2.,
-                      lambda xf: xf - x)
+                      x + 2.)
     expected = """{ lambda ; a:f32[]. let
     b:bool[] = ge a 0.0:f32[]
     c:f32[] = add a 1.0:f32[]
@@ -7941,10 +7941,10 @@ def f(c, x):
     jax.lax.scan(f, 0, jnp.arange(4))
 
   def test_cond_traceback(self):
-    if sys.version_info < (3, 14):
+    if sys.version_info < (3, 13):
       # Fails because 3.11 adds an extra stack frame due to a list comprehension
       self.skipTest("Expected failure.")
-    expected_depth = 8
+    expected_depth = 4
     init_depth = self.cur_depth()
 
     def f():

tests/core_test.py

@@ -448,15 +448,11 @@ class JaxprTypeChecks(jtu.JaxTestCase):
 
   def setUp(self):
     super().setUp()
-    lax_control_flow._initial_style_open_jaxpr.cache_clear()
-    lax_control_flow._initial_style_jaxpr.cache_clear()
     lax_control_flow.common._dedup_consts.cache_clear()
     lax_control_flow.common._pad_constvars.cache_clear()
 
   def tearDown(self):
     super().tearDown()
-    lax_control_flow._initial_style_open_jaxpr.cache_clear()
-    lax_control_flow._initial_style_jaxpr.cache_clear()
     lax_control_flow.common._dedup_consts.cache_clear()
     lax_control_flow.common._pad_constvars.cache_clear()