Repository: google/jax

Issue:
[FR] support batching rule for `while`
Currently, `vmap` seems do not support a function which involves `lax.while_loop`. For example, the following script will throw `NotImplementedError: Batching rule for 'while' not implemented`:
```
import jax.numpy as np
from jax import jit, lax, vmap

@jit
def f(x):
    return lax.while_loop(lambda x: x <= 0, lambda x: x + 1, x)

g = vmap(f)
y = g(np.arange(3.))
```

In case it is complicated to support `while_loop` for `vmap`, then could someone let me know other alternatives? I know some other alternatives from *gufuncs* tutorial such as using `onp.vectorize` or python for loop but they are slow. To make `g` fast, can I write `f` in C/Cython code and use vmap with it? Or should I write the whole `g` in C/Cython code? I have little experience with C/Cython but will try to learn if that is the only option.

I just want to use `g` as a primitive function and will define `jvp_rule` for it separately.

Thank you for any help in advance!


Discussion:
Thanks for requesting this!

No need for C/Cython; we can add a batching rule for `lax.while_loop`. In general there are some rules we haven't gotten to implementing yet, and we tend to wait for users to ask for them as a natural way to prioritize our work.
Thank you @mattjj ! That's would be a nice feature. I have been playing with JAX for a few days and really love it.
Is it correct to assume that a batched `while_loop` would still require a `cond_fun` with scalar boolean output or do you intend to support a batched `cond_fun` that outputs a boolean for each sample in the batch for which the loop should be continued?
That’s a great question! I was thinking the latter, but I haven’t thought through it all the way. If we wanted the whole loop to be vectorized then we should turn the condition into an `np.all` check along the batch dimension and use a mask to keep updating only those loops that are still running. 

Thoughts?
That would be great!
This is one of the reasons that I wasn't able to disentangle vectorization and masking in Matchbox; it might be reasonable to allow vmap of while (where the condition contains a vmapped dimension) to invoke the masking interpreter when that's ready.
@jekbradbury Hmm perhaps, but in this case couldn't the masking stay at the level of the while loop primitive, rather than being pushed down into the body function (like a JAX interpreter would do)?

Resolving patch:
diff --git a/jax/interpreters/batching.py b/jax/interpreters/batching.py
--- a/jax/interpreters/batching.py
+++ b/jax/interpreters/batching.py
@@ -35,28 +35,27 @@
 map = safe_map
 
 
-def batch(fun, in_vals, in_dims, out_dim_target):
+def batch(fun, in_vals, in_dims, out_dim_dst):
   sizes = reduce(set.union, map(dimsize, in_dims, in_vals))
   if not sizes:
     return fun.call_wrapped(*in_vals), None  # no mapped dimensions
   elif len(sizes) == 1:
-    out_val, out_dim = batch_transform(fun).call_wrapped(in_vals, in_dims)
-    return moveaxis(sizes.pop(), out_dim_target, out_dim, out_val)
+    sz = sizes.pop()
+    return batch_transform(fun, sz, in_dims, out_dim_dst).call_wrapped(in_vals)
   else:
     raise TypeError("got inconsistent map dimension sizes: {}".format(sizes))
 
 
-# TODO(mattjj,dougalm): could call batch_subtrace here (a bit redundant)
 @transformation
-def batch_transform(vals, dims):
+def batch_transform(size, in_dims, out_dim_dst, vals):
   with new_master(BatchTrace) as master:
     trace = BatchTrace(master, core.cur_sublevel())
-    in_tracers = map(partial(BatchTracer, trace), vals, dims)
+    in_tracers = map(partial(BatchTracer, trace), vals, in_dims)
     out_tracer = yield in_tracers
     out_tracer = trace.full_raise(out_tracer)
     out_val, out_dim = out_tracer.val, out_tracer.batch_dim
     del master
-  yield (out_val, out_dim)
+  yield moveaxis(size, out_dim_dst, out_dim, out_val)
 
 
 @transformation_with_aux
@@ -308,13 +307,21 @@ def move_dim_to_front(x, dim):
 def dimsize(dim, x):
   aval = get_aval(x)
   if type(aval) is AbstractTuple:
-    return reduce(set.union, map(partial(dimsize, dim), x))
-  elif type(dim) is int:
-    return {x.shape[dim]}
-  elif dim is None:
-    return set()
+    if type(dim) is tuple:
+      return reduce(set.union, map(dimsize, dim, x))
+    elif type(dim) is int:
+      return reduce(set.union, map(partial(dimsize, dim), x))
+    elif dim is None:
+      return set()
+    else:
+      raise TypeError(type(dim))
   else:
-    raise TypeError(type(dim))
+    if type(dim) is int:
+      return {x.shape[dim]}
+    elif dim is None:
+      return set()
+    else:
+      raise TypeError(type(dim))
 
 def moveaxis(sz, dst, src, x):
   aval = get_aval(x)
diff --git a/jax/lax.py b/jax/lax.py
--- a/jax/lax.py
+++ b/jax/lax.py
@@ -3665,10 +3665,56 @@ def _while_loop_translation_rule(c, init_val, cond_consts, body_consts,
   full_ans = c.While(cond_computation, body_computation, loop_carry)
   return c.GetTupleElement(full_ans, 0)
 
+def _while_loop_batching_rule(batched_args, batch_dims, aval_out, cond_jaxpr,
+                              body_jaxpr):
+  # See https://github.com/google/jax/issues/441 for a discussion.
+  # To batch a while_loop, we need to do some masking, since the elements of the
+  # batch may run for different numbers of iterations. We perform that masking
+  # using lax.select, and keep the loop running so long as any of the batch
+  # elements need by effectively using an np.any(...) in the cond_fun.
+  # The basic strategy here is to lift `cond_jaxpr` and `body_jaxpr` back into
+  # traceable Python functions using `core.eval_jaxpr`. Then we can batch them
+  # using `batching.batch_transform` (the transform underlying `api.vmap`). This
+  # code also avoids broadcasting `cond_consts` and `body_consts`.
+  init_val, cond_consts, body_consts = batched_args
+  init_val_bd, cond_consts_bd, body_consts_bd = batch_dims
+
+  sizes = _reduce(set.union, map(batching.dimsize, batch_dims, batched_args))
+  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
+
+  def batched_cond_fun(batched_loop_carry):
+    @lu.wrap_init
+    def lifted(loop_carry, cond_consts):
+      return core.eval_jaxpr(cond_jaxpr, cond_consts, (), loop_carry)
+    f = batching.batch_transform(lifted, size, (init_val_bd, cond_consts_bd), 0)
+    preds = f.call_wrapped((batched_loop_carry, cond_consts))
+    return reduce(preds, onp.array(False), bitwise_or, [0])
+
+  def batched_body_fun(batched_loop_carry):
+    @lu.wrap_init
+    def lifted(loop_carry, cond_consts, body_consts):
+      pred = core.eval_jaxpr(cond_jaxpr, cond_consts, (), loop_carry)
+      new_loop_carry = core.eval_jaxpr(body_jaxpr, body_consts, (), loop_carry)
+      return select(pred, new_loop_carry, loop_carry)
+    f = batching.batch_transform(
+        lifted, size, (init_val_bd, cond_consts_bd, body_consts_bd), init_val_bd)
+    return f.call_wrapped((batched_loop_carry, cond_consts, body_consts))
+
+  return while_loop(batched_cond_fun, batched_body_fun, init_val), init_val_bd
+
 while_p = Primitive('while')
 while_p.def_impl(partial(xla.apply_primitive, while_p))
 while_p.def_abstract_eval(_while_loop_abstract_eval)
 xla.translations[while_p] = _while_loop_translation_rule
+batching.primitive_batchers[while_p] = _while_loop_batching_rule
 
 def _unpack_eqn(invar, outvars):
   return core.JaxprEqn([invar], outvars, core.identity_p, (), True, {})


Test patch:
diff --git a/tests/batching_test.py b/tests/batching_test.py
--- a/tests/batching_test.py
+++ b/tests/batching_test.py
@@ -833,6 +833,35 @@ def f(R):
 
     H = hessian(f)(R)  # don't crash on UnshapedArray
 
+  def testWhileLoop(self):
+    def fun(x):
+      return lax.while_loop(lambda x: x < 3, lambda x: x + 2, x)
+
+    ans = vmap(fun)(onp.array([0, 1, 2, 3]))
+    expected = onp.array([4, 3, 4, 3])
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+    fun = jit(fun)
+    ans = vmap(fun)(onp.array([0, 1, 2, 3]))
+    expected = onp.array([4, 3, 4, 3])
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+  def testWhileLoopCondConstsBatched(self):
+    def fun(x, y):
+      return lax.while_loop(lambda x: x < y, lambda x: x + 2, x)
+
+    ans = vmap(fun, in_axes=(None, 0))(0, onp.array([2, 3]))
+    expected = onp.array([2, 4])
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+  def testWhileLoopBodyConstsBatched(self):
+    def fun(x, y):
+      return lax.while_loop(lambda x: x < 3, lambda x: x + y, x)
+
+    ans = vmap(fun, in_axes=(None, 0))(0, onp.array([2, 3]))
+    expected = onp.array([4, 3])
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
 
 if __name__ == '__main__':
   absltest.main()