Repository: google/jax

Issue:
lax.cond leads to excessive compilation in JAX 0.4.1
Reported in https://github.com/google/jax/discussions/14032#discussioncomment-4718316. Minimal reproduction:
```python
import jax
import jaxlib
from jax import lax

print(jaxlib.__version__)
print(jax.__version__)

jax.config.update('jax_log_compiles', True)

def f(x):
  return x + 1

def g(x):
  return x + 2

for x in range(10):
  lax.cond(x, f, g, x)
```
Output in JAX v0.3.25:
```
0.3.25
0.3.25
WARNING:jax._src.dispatch:Finished tracing + transforming prim_fun for jit in 0.0018656253814697266 sec
WARNING:jax._src.dispatch:Compiling prim_fun (139732125277056 for args (ShapedArray(int32[]), ShapedArray(int32[], weak_type=True)).
WARNING:jax._src.dispatch:Finished XLA compilation of cond in 0.013795614242553711 sec
```
Output in JAX v0.4.1:
```
0.4.1
0.4.1
WARNING:jax._src.dispatch:Finished tracing + transforming jit(cond) in 0.0008955001831054688 sec
WARNING:jax.interpreters.pxla:Compiling prim_fun (139864230252032) for with global shapes and types (ShapedArray(int32[]), ShapedArray(int32[], weak_type=True)). Argument mapping: (OpShardingSharding({replicated}), OpShardingSharding({replicated})).
WARNING:jax._src.dispatch:Finished XLA compilation of jit(cond) in 0.02367258071899414 sec
WARNING:jax._src.dispatch:Finished tracing + transforming jit(cond) in 0.0006630420684814453 sec
WARNING:jax.interpreters.pxla:Compiling prim_fun (139864230252112) for with global shapes and types (ShapedArray(int32[]), ShapedArray(int32[], weak_type=True)). Argument mapping: (OpShardingSharding({replicated}), OpShardingSharding({replicated})).
WARNING:jax._src.dispatch:Finished XLA compilation of jit(cond) in 0.023768901824951172 sec
WARNING:jax._src.dispatch:Finished tracing + transforming jit(cond) in 0.0007023811340332031 sec
WARNING:jax.interpreters.pxla:Compiling prim_fun (139864230322960) for with global shapes and types (ShapedArray(int32[]), ShapedArray(int32[], weak_type=True)). Argument mapping: (OpShardingSharding({replicated}), OpShardingSharding({replicated})).
WARNING:jax._src.dispatch:Finished XLA compilation of jit(cond) in 0.018657445907592773 sec
WARNING:jax._src.dispatch:Finished tracing + transforming jit(cond) in 0.000583648681640625 sec
WARNING:jax.interpreters.pxla:Compiling prim_fun (139864230251872) for with global shapes and types (ShapedArray(int32[]), ShapedArray(int32[], weak_type=True)). Argument mapping: (OpShardingSharding({replicated}), OpShardingSharding({replicated})).
WARNING:jax._src.dispatch:Finished XLA compilation of jit(cond) in 0.0319209098815918 sec
WARNING:jax._src.dispatch:Finished tracing + transforming jit(cond) in 0.01588129997253418 sec
WARNING:jax.interpreters.pxla:Compiling prim_fun (139864230325360) for with global shapes and types (ShapedArray(int32[]), ShapedArray(int32[], weak_type=True)). Argument mapping: (OpShardingSharding({replicated}), OpShardingSharding({replicated})).
WARNING:jax._src.dispatch:Finished XLA compilation of jit(cond) in 0.020395517349243164 sec
WARNING:jax._src.dispatch:Finished tracing + transforming jit(cond) in 0.0006628036499023438 sec
WARNING:jax.interpreters.pxla:Compiling prim_fun (139864230323680) for with global shapes and types (ShapedArray(int32[]), ShapedArray(int32[], weak_type=True)). Argument mapping: (OpShardingSharding({replicated}), OpShardingSharding({replicated})).
WARNING:jax._src.dispatch:Finished XLA compilation of jit(cond) in 0.018640756607055664 sec
WARNING:jax._src.dispatch:Finished tracing + transforming jit(cond) in 0.0006759166717529297 sec
WARNING:jax.interpreters.pxla:Compiling prim_fun (139864230322400) for with global shapes and types (ShapedArray(int32[]), ShapedArray(int32[], weak_type=True)). Argument mapping: (OpShardingSharding({replicated}), OpShardingSharding({replicated})).
WARNING:jax._src.dispatch:Finished XLA compilation of jit(cond) in 0.017685890197753906 sec
WARNING:jax._src.dispatch:Finished tracing + transforming jit(cond) in 0.0006914138793945312 sec
WARNING:jax.interpreters.pxla:Compiling prim_fun (139864230369536) for with global shapes and types (ShapedArray(int32[]), ShapedArray(int32[], weak_type=True)). Argument mapping: (OpShardingSharding({replicated}), OpShardingSharding({replicated})).
WARNING:jax._src.dispatch:Finished XLA compilation of jit(cond) in 0.018553495407104492 sec
WARNING:jax._src.dispatch:Finished tracing + transforming jit(cond) in 0.0007042884826660156 sec
WARNING:jax.interpreters.pxla:Compiling prim_fun (139864230370576) for with global shapes and types (ShapedArray(int32[]), ShapedArray(int32[], weak_type=True)). Argument mapping: (OpShardingSharding({replicated}), OpShardingSharding({replicated})).
WARNING:jax._src.dispatch:Finished XLA compilation of jit(cond) in 0.03286623954772949 sec
WARNING:jax._src.dispatch:Finished tracing + transforming jit(cond) in 0.0027599334716796875 sec
WARNING:jax.interpreters.pxla:Compiling prim_fun (139864230367856) for with global shapes and types (ShapedArray(int32[]), ShapedArray(int32[], weak_type=True)). Argument mapping: (OpShardingSharding({replicated}), OpShardingSharding({replicated})).
WARNING:jax._src.dispatch:Finished XLA compilation of jit(cond) in 0.02003955841064453 sec
```
Assigning to @yashk2810 because it's possible this is somehow related to the jax.Array change


Discussion:
Actually, I just checked and `config.jax_array` has no effect on this.
Maybe try setting `jax_experimental_subjaxpr_lowering_cache` config to True and see what happens?
`jax_experimental_subjaxpr_lowering_cache` doesn't seem to affect it either way.
The root cause seems to be #12877, which added these lines https://github.com/google/jax/blob/a37121e19512ea5ee0ad523eda39fa5bbd8c5442/jax/_src/lax/control_flow/conditionals.py#L258-L261

Commenting them out fixes the issue.

Resolving patch:
diff --git a/jax/_src/core.py b/jax/_src/core.py
--- a/jax/_src/core.py
+++ b/jax/_src/core.py
@@ -26,7 +26,7 @@
 from operator import attrgetter
 import threading
 import types
-from typing import (Any, Callable, ClassVar, DefaultDict, Dict,
+from typing import (Any, Callable, ClassVar, DefaultDict, Dict, FrozenSet,
                     Generator, Hashable, Iterable, Iterator, List,
                     NamedTuple, Optional, Sequence, Set, Tuple, Type,
                     Union, cast)
@@ -2444,6 +2444,13 @@ def subst_axis_names_jaxpr(jaxpr: Union[Jaxpr, ClosedJaxpr], subst: AxisSubst):
     return jaxpr
   return do_subst_axis_names_jaxpr(jaxpr, subst)
 
+def replace_jaxpr_effects(jaxpr: ClosedJaxpr, effects: Effects):
+  return _replace_jaxpr_effects(jaxpr, frozenset(effects))
+
+@weakref_lru_cache
+def _replace_jaxpr_effects(jaxpr: ClosedJaxpr, effects: FrozenSet[Effect]):
+  return jaxpr.replace(jaxpr=jaxpr.jaxpr.replace(effects=set(effects)))
+
 
 axis_substitution_rules: Dict[Primitive, Callable[[ParamDict, AxisSubst, bool], ParamDict]] = {}
 
diff --git a/jax/_src/lax/control_flow/conditionals.py b/jax/_src/lax/control_flow/conditionals.py
--- a/jax/_src/lax/control_flow/conditionals.py
+++ b/jax/_src/lax/control_flow/conditionals.py
@@ -32,6 +32,7 @@
 from jax.interpreters import xla
 from jax.tree_util import tree_flatten, tree_unflatten
 from jax._src import ad_util
+from jax._src.core import replace_jaxpr_effects
 from jax._src import dtypes
 from jax._src import source_info_util
 from jax._src import util
@@ -255,10 +256,8 @@ def cond(pred, true_fun, false_fun, *operands):
     # Raise index in case of effects to allow data-dependence-based discharging
     # of those effects (even if they don't have an explicit data dependence).
     index = core.raise_as_much_as_possible(index)
-  false_jaxpr = false_jaxpr.replace(
-      jaxpr=false_jaxpr.jaxpr.replace(effects=joined_effects))
-  true_jaxpr = true_jaxpr.replace(
-      jaxpr=true_jaxpr.jaxpr.replace(effects=joined_effects))
+  false_jaxpr = replace_jaxpr_effects(false_jaxpr, joined_effects)
+  true_jaxpr = replace_jaxpr_effects(true_jaxpr, joined_effects)
 
   linear = [False] * len(consts) + linear_ops
   out = cond_p.bind(


Test patch:
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
--- a/tests/lax_control_flow_test.py
+++ b/tests/lax_control_flow_test.py
@@ -2400,6 +2400,21 @@ def f_jvp(primals, tangents):
     x = np.arange(3, dtype='float32')
     jax.jvp(g, (x,), (x,))  # doesn't crash
 
+  def test_cond_excessive_compilation(self):
+    # Regression test for https://github.com/google/jax/issues/14058
+    def f(x):
+      return x + 1
+
+    def g(x):
+      return x + 2
+
+    with jtu.count_jit_and_pmap_compiles() as count:
+      for x in range(10):
+        lax.cond(x, f, g, x)
+    # Should observe a maximum of 4 compiles: convert_element_type, f, g, cond
+    # In #14058, this was observed to be 31 compiles.
+    self.assertLess(count[0], 5)
+
   @parameterized.named_parameters(
       {"testcase_name": f"_dtype={dtype.__name__}", "dtype": dtype}
       for dtype in jtu.dtypes.all_integer)