Repository: google/jax

Issue:
While-loop vmap bug
The following code block runs as expected on jax 0.1.63, jaxlib 0.1.45, but fails on all later versions, including master:

```python
import jax
import jax.numpy as np
from jax.experimental import loops

def test(a,b):
    with loops.Scope() as s:
        s.val = 0
        s.i = 0
        s.j = 0
        for _ in s.while_range(lambda: s.i < a + 1):
            s.j = 0
            for _ in s.while_range(lambda: s.j < b + 1):
                s.val += s.i + s.j
                s.j += 1
            s.i += 1
        return s.val

# vectorized version
vmap_test = jax.vmap(test, (0,0))
arr = np.arange(5)
vmap_test(arr, arr)
```
<details><summary>Click for Traceback</summary>
<p>

```python
Traceback (most recent call last):
  File "test.py", line 21, in <module>
    print(vmap_test(arr, arr))
  File "/home/adabbott/Git/jax/jax/jax/api.py", line 858, in batched_fun
    lambda: flatten_axes(out_tree(), out_axes))
  File "/home/adabbott/Git/jax/jax/jax/interpreters/batching.py", line 34, in batch
    return batched_fun.call_wrapped(*in_vals)
  File "/home/adabbott/Git/jax/jax/jax/linear_util.py", line 150, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "test.py", line 12, in test
    for _ in s.while_range(lambda: s.j < b + 1):
  File "/home/adabbott/Git/jax/jax/jax/experimental/loops.py", line 341, in __next__
    self.end_tracing_body()
  File "/home/adabbott/Git/jax/jax/jax/experimental/loops.py", line 407, in end_tracing_body
    carried_init_vals, body_typed_jaxpr, body_const_vals)
  File "/home/adabbott/Git/jax/jax/jax/experimental/loops.py", line 576, in build_output_vals
    body_jaxpr=body_typed_jaxpr)
  File "/home/adabbott/Git/jax/jax/jax/core.py", line 212, in bind
    out_tracer = top_trace.process_primitive(self, tracers, kwargs)
  File "/home/adabbott/Git/jax/jax/jax/interpreters/partial_eval.py", line 141, in process_primitive
    return custom_partial_eval_rules[primitive](self, *tracers, **params)
  File "/home/adabbott/Git/jax/jax/jax/lax/lax_control_flow.py", line 517, in _while_partial_eval
    body_jaxpr=body_jaxpr_known)
  File "/home/adabbott/Git/jax/jax/jax/core.py", line 212, in bind
    out_tracer = top_trace.process_primitive(self, tracers, kwargs)
  File "/home/adabbott/Git/jax/jax/jax/interpreters/batching.py", line 134, in process_primitive
    val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
  File "/home/adabbott/Git/jax/jax/jax/lax/lax_control_flow.py", line 391, in _while_loop_batching_rule
    body_nconsts=body_nconsts, body_jaxpr=body_jaxpr_batched)
  File "/home/adabbott/Git/jax/jax/jax/core.py", line 209, in bind
    return self.impl(*args, **kwargs)
  File "/home/adabbott/Git/jax/jax/jax/interpreters/xla.py", line 217, in apply_primitive
    compiled_fun = xla_primitive_callable(prim, *map(arg_spec, args), **params)
  File "/home/adabbott/Git/jax/jax/jax/interpreters/xla.py", line 248, in xla_primitive_callable
    *avals, **params)
  File "/home/adabbott/Git/jax/jax/jax/interpreters/xla.py", line 295, in primitive_computation
    *xla_args, **params)
  File "/home/adabbott/Git/jax/jax/jax/lax/lax_control_flow.py", line 332, in _while_loop_translation_rule
    new_z = _map(partial(_pred_bcast_select, body_c, body_pred), new_z, z)
  File "/home/adabbott/Git/jax/jax/jax/util.py", line 34, in safe_map
    return list(map(f, *args))
  File "/home/adabbott/Git/jax/jax/jax/lax/lax_control_flow.py", line 350, in _pred_bcast_select
    assert pred_shape == x_shape[:len(pred_shape)] == y_shape[:len(pred_shape)]
AssertionError
```

</p>
</details>

It appears to only occur when the nested while-loop variable `b` is vectorized:

```python
# this works
vmap_test = jax.vmap(test, (0,None))
vmap_test(arr, 3)

# this fails
vmap_test = jax.vmap(test, (None,0))
vmap_test(3, arr)
```



Discussion:
I should also note the same behavior occurs when using `jax.lax.while_loop` directly, without the convenience of the `loops` module:
<details><summary>Click for pure jax.lax.while_loop version </summary>
<p>

```python
import jax
import jax.numpy as np

def test(a,b):
    val = 0
    i = 0
    j = 0

    condfun_1 = lambda inp: inp[1] < a + 1 
    condfun_2 = lambda inp: inp[2] < b + 1 

    def bodyfun_1(inp):
        val, i, j = inp
        j = 0
        def bodyfun_2(inp):
            val, i, j = inp
            val += i + j
            j += 1
            return (val, i, j)
        result = jax.lax.while_loop(condfun_2, bodyfun_2, (val,i,j))
        val = result[0]
        i += 1
        return (val, i, j)
    result = jax.lax.while_loop(condfun_1, bodyfun_1, (val,i,j))
    return result[0]

arr = np.arange(5)
vmap_test = jax.vmap(test, (0,0))
vmap_test(arr, arr)
```

</p>
</details>

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
@@ -329,7 +329,7 @@ def _while_loop_translation_rule(c, axis_env, name_stack, avals, backend, *args,
     body_pred, = xla.jaxpr_subcomp(body_c, cond_jaxpr.jaxpr, backend, axis_env,
                                    _map(partial(xb.constant, body_c), cond_jaxpr.literals),
                                    extend_name_stack(name_stack, 'body_pred'), *(x + z))
-    new_z = _map(partial(_pred_bcast_select, body_c, body_pred), new_z, z)
+    new_z = _map(partial(_pred_bcast_select, body_c, body_pred), new_z, z, body_jaxpr.out_avals)
     assert _map(body_c.get_shape, new_z) == _map(body_c.get_shape, z) # no broadcast
   new_carry = xops.Tuple(body_c, list(itertools.chain(x, y, new_z)))
 
@@ -338,13 +338,14 @@ def _while_loop_translation_rule(c, axis_env, name_stack, avals, backend, *args,
   _,  _, z = split_list(ans_elts, [cond_nconsts, body_nconsts])
   return xops.Tuple(c, z)
 
-def _pred_bcast_select(c, pred, x, y):
+def _pred_bcast_select(c, pred, x, y, x_y_aval: core.AbstractValue):
   pred_shape = c.get_shape(pred).dimensions()
   x_shape = c.get_shape(x).dimensions()
   y_shape = c.get_shape(y).dimensions()
   assert x_shape == y_shape
-  if not c.get_shape(x).is_array() and not c.get_shape(y).is_array():
-    # Two tokens
+  if x_y_aval is core.abstract_unit:
+    return x
+  elif x_y_aval is core.abstract_token:
     return xops.AfterAll(c, [x, y])
   else:
     assert pred_shape == x_shape[:len(pred_shape)] == y_shape[:len(pred_shape)]


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
@@ -334,6 +334,38 @@ def fun(x, y):
     expected = (np.array([4, 3]), np.array([1, 2]))
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  def test_issue_3204(self):
+    # Error during XLA code generation for vmap of nested loops
+    def test(a, b):
+      val = 0
+      i = 0
+      j = 0
+
+      condfun_1 = lambda inp: inp[1] < a + 1
+      condfun_2 = lambda inp: inp[2] < b + 1
+
+      def bodyfun_1(inp):
+        val, i, j = inp
+        j = 0
+
+        def bodyfun_2(inp):
+          val, i, j = inp
+          val += i + j
+          j += 1
+          return (val, i, j)
+
+        result = lax.while_loop(condfun_2, bodyfun_2, (val, i, j))
+        val = result[0]
+        i += 1
+        return (val, i, j)
+
+      result = lax.while_loop(condfun_1, bodyfun_1, (val, i, j))
+      return result[0]
+
+    arr = np.arange(5)
+    vmap_test = api.vmap(test, (0, 0))
+    vmap_test(arr, arr)
+
   def testForiLoopErrors(self):
     """Test typing error messages for while."""
     with self.assertRaisesRegex(