Repository: google/jax
Issue: while_loop vmap error

This one is reported by @fehiepsi in [another thread](https://github.com/google/jax/pull/486#issuecomment-471140362). Quoting the repro from there:

```
# generate a random number in the interval [0, 0.5]
def f(key):
    def body_fn(uk):
        key = uk[1]
        u = random.uniform(key, ())
        key, _ = random.split(key)
        return u, key

    u, _ = lax.while_loop(lambda uk: uk[0] > 0.5, body_fn, (1., key))
    #u = random.uniform(key, ())  # this is fine
    return u

print(f(random.PRNGKey(0)))  # no error
print(vmap(f)(random.split(random.PRNGKey(0), 2)))  # TypeError: 'NoneType' object cannot be interpreted as an integer
```

Changed files:
jax/lax.py, jax/interpreters/batching.py, tests/batching_test.py

diff --git a/jax/interpreters/batching.py b/jax/interpreters/batching.py
index 6578176a3..529fec143 100644
--- a/jax/interpreters/batching.py
+++ b/jax/interpreters/batching.py
@@ -285,24 +285,11 @@ primitive_batchers[zeros_like_p] = zeros_like_batched
 
 
 def bdim_at_front(x, bdim, broadcast_size=1, force_broadcast=False):
-  if bdim is None:
-    return broadcast(x, broadcast_size, force_broadcast=force_broadcast)
-  else:
-    return move_dim_to_front(x, bdim)
+  return moveaxis(broadcast_size, 0, bdim, x, force_broadcast=force_broadcast)
 
 def move_dim_to_front(x, dim):
-  aval = get_aval(x)
-  if type(aval) is AbstractTuple:
-    return pack(map(partial(move_dim_to_front, dim=dim), x))
-  elif isinstance(aval, ShapedArray):
-    assert 0 <= dim < onp.ndim(x)
-    if dim == 0:
-      return x
-    else:
-      perm = (dim,) + tuple(range(dim)) + tuple(range(dim + 1, onp.ndim(x)))
-      return x.transpose(perm)
-  else:
-    raise TypeError(type(x))
+  assert dim is not None
+  return moveaxis(None, 0, dim, x)
 
 def dimsize(dim, x):
   aval = get_aval(x)
@@ -323,7 +310,7 @@ def dimsize(dim, x):
     else:
       raise TypeError(type(dim))
 
-def moveaxis(sz, dst, src, x):
+def moveaxis(sz, dst, src, x, force_broadcast=True):
   aval = get_aval(x)
   if type(aval) is AbstractTuple:
     if type(src) is tuple and type(dst) is tuple:
@@ -341,7 +328,7 @@ def moveaxis(sz, dst, src, x):
       return x
     else:
       if src is None:
-        x = broadcast(x, sz, force_broadcast=True)
+        x = broadcast(x, sz, force_broadcast=force_broadcast)
         src = 0
         dst_ = dst % (aval.ndim + 1)
       if src == dst_:
diff --git a/jax/lax.py b/jax/lax.py
index 9022ab8c6..f510e1854 100644
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -3683,12 +3683,11 @@ def _while_loop_batching_rule(batched_args, batch_dims, aval_out, cond_jaxpr,
   size = sizes.pop()
   assert not sizes
 
-  if init_val_bd is None:
-    # TODO(mattjj): if cond_consts_bd is also None, we could keep cond_fun
-    # unbatched and avoid the masking logic, but we ignore that optimiztaion
-    init_val = batching.bdim_at_front(init_val, init_val_bd, size,
-                                      force_broadcast=True)
-    init_val_bd = 0
+  # TODO(mattjj): if cond_consts_bd is also None, we could keep cond_fun
+  # unbatched and avoid the masking logic, but we ignore that optimiztaion
+  init_val = batching.bdim_at_front(init_val, init_val_bd, size,
+                                    force_broadcast=True)
+  init_val_bd = 0
 
   def batched_cond_fun(batched_loop_carry):
     @lu.wrap_init
diff --git a/tests/batching_test.py b/tests/batching_test.py
index 2c1c26f99..1fea08e39 100644
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -913,6 +913,19 @@ class BatchingTest(jtu.JaxTestCase):
     expected = (onp.array([10, 11]), onp.array([20, 20]))
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  def testIssue489(self):
+    def f(key):
+      def body_fn(uk):
+        key = uk[1]
+        u = random.uniform(key, (), dtype=np.float64)
+        key, _ = random.split(key)
+        return u, key
+
+      u, _ = lax.while_loop(lambda uk: uk[0] > 0.5, body_fn, (1., key))
+      return u
+
+    print(vmap(f)(random.split(random.PRNGKey(0), 2)))  # no crash
+
 
 if __name__ == '__main__':
   absltest.main()