Repository: google/jax

Issue:
Bug with jax.disable_jit (or lax.cond, or vmap?)
Attempt at minimal reproducing example: 
```python
def fn(t):
  return jax.lax.cond(t > 0, 0, lambda x: 0, 0, lambda x: 1)
fn = jax.vmap(fn)

with jax.disable_jit():
  outputs = fn(jnp.array([1]))
```
fails with 
```
Encountered value: Traced<ShapedArray(bool[])>with<BatchTrace(level=11/0)>
  with val = DeviceArray([ True], dtype=bool)
       batch_dim = 0
```
It does not fail without the `disable_jit()` and it does not fail without the `vmap` (input a scalar in that case).
disable_jit extensions to control flow constructs break with vmap
The extension of `jax.disable_jit()` context to the control flow elements (e.g. `while_loop`, `cond`) isn't correct in the presence of batch tracers for `vmap`.   `while_loop` and `cond` are rewritten to use python control flow under this context, but conditional branching can't be evaluated on the batchtracers present when functions using these control flow constructs are inside a `vmap`.

e.g.
```python
def trivial_while(y):
  return lax.while_loop(lambda x: x < 10.0, lambda x: x + 1.0, y)
# works:
jax.vmap(trivial_while)(jnp.array([3.0,4.0]))
# throws a ConcretizationTypeError:
with jax.disable_jit():
  jax.vmap(trivial_while)(jnp.array([3.0,4.0]))
```
likewise:
```python
def trivial_cond(x):
  return lax.cond(x < 1.0, x, lambda x: x, x, lambda x: 2*x)
# works
jax.vmap(trivial_cond)(jnp.array([0.5,2.0]))
# throws a ConcretizationTypeError:
with jax.disable_jit():
  jax.vmap(trivial_cond)(jnp.array([0.5,2.0]))
```

Perhaps the disabled_jit python implementations for these control flow constructs could peak at the tracer value's `val` field which I suspect will be populated in these use-cases?


Discussion:
Thanks for raising this! This is a known problem, though I forget if there's already an issue for it.

The problem is that:
1. when in a `disable_jit` context, the control flow primitives like `cond` [just fall back to a Python implementation](https://github.com/google/jax/blob/7c687b245b34397c13563a714ad9bf0290b419e3/jax/lax/lax_control_flow.py#L530-L534) motivated by better debugging, yet
2. this Python implementation isn't `vmap`-traceable (because in this case `vmap` abstracts the boolean `t > 0` to the Shaped level).

Probably the best option is to run the Python implementation only when `bool(pred)` can be evaluated at trace time (i.e. when that call doesn't raise an abstract value error), and call into the non-Python version when it fails.
#2598 is possibly related? In particular commenting out [these lines](https://github.com/google/jax/blob/2c4ced2143375ed8a4f8e467c501d9c794d385cb/jax/lax/lax_control_flow.py#L219-L223) might fix it?

Resolving patch:
diff --git a/jax/lax/lax_control_flow.py b/jax/lax/lax_control_flow.py
--- a/jax/lax/lax_control_flow.py
+++ b/jax/lax/lax_control_flow.py
@@ -33,7 +33,7 @@
 from jax import util
 from jax.lax import lax
 from jax import linear_util as lu
-from jax.abstract_arrays import ShapedArray, raise_to_shaped
+from jax.abstract_arrays import ConcreteArray, ShapedArray, raise_to_shaped
 from jax.api_util import flatten_fun_nokwargs, apply_flat_fun_nokwargs
 from jax.core import get_aval
 from jax.interpreters import ad
@@ -267,10 +267,15 @@ def while_loop(cond_fun, body_fun, init_val):
     The output from the final iteration of body_fun, of type ``a``.
   """
   if jax.api._jit_is_disabled():
-    val = init_val
-    while cond_fun(val):
-      val = body_fun(val)
-    return val
+    try:
+      val = init_val
+      while cond_fun(val):
+        val = body_fun(val)
+      return val
+    except core.ConcretizationTypeError:
+      # Can't run this while_loop in Python (e.g. because there's a vmap
+      # transformation on it), so we fall back to the primitive version.
+      pass
 
   init_vals, in_tree = tree_flatten((init_val,))
   init_avals = tuple(_map(_abstractify, init_vals))
@@ -593,7 +598,7 @@ def _cond(pred, true_fun: Callable, false_fun: Callable, operand):
       msg = ("Pred type must be either boolean or number, got {}.")
       raise TypeError(msg.format(pred_dtype))
 
-  if jax.api._jit_is_disabled():
+  if jax.api._jit_is_disabled() and isinstance(core.get_aval(pred), ConcreteArray):
     if pred:
       return true_fun(operand)
     else:


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
@@ -1966,6 +1966,22 @@ def cumsum(x, reverse):
     with api.disable_jit():
       self.assertAllClose(np.cumsum(x[::-1])[::-1], cumsum(x, True), check_dtypes=False)
 
+  def test_disable_jit_cond_with_vmap(self):
+    # https://github.com/google/jax/issues/3093
+    def fn(t):
+      return lax.cond(t > 0, 0, lambda x: 0, 0, lambda x: 1)
+    fn = api.vmap(fn)
+
+    with api.disable_jit():
+      outputs = fn(jnp.array([1]))  # doesn't crash
+
+  def test_disable_jit_while_loop_with_vmap(self):
+    # https://github.com/google/jax/issues/2823
+    def trivial_while(y):
+      return lax.while_loop(lambda x: x < 10.0, lambda x: x + 1.0, y)
+    with api.disable_jit():
+      api.vmap(trivial_while)(jnp.array([3.0,4.0]))  # doesn't crash
+
 
 if __name__ == '__main__':
   absltest.main()
diff --git a/tests/ode_test.py b/tests/ode_test.py
--- a/tests/ode_test.py
+++ b/tests/ode_test.py
@@ -155,6 +155,14 @@ def g(x):
     rtol = {np.float64: 2e-15}
     self.assertAllClose(ans, expected, check_dtypes=False, atol=atol, rtol=rtol)
 
+  def test_disable_jit_odeint_with_vmap(self):
+    # https://github.com/google/jax/issues/2598
+    with jax.disable_jit():
+      t = jax.numpy.array([0.0, 1.0])
+      x0_eval = jax.numpy.zeros((5, 2))
+      f = lambda x0: odeint(lambda x, _t: x, x0, t)
+      jax.vmap(f)(x0_eval)  # doesn't crash
+
 
 if __name__ == '__main__':
   absltest.main()