Repository: google/jax

Issue:
add cond batching rule
Covers both the predication (batched condition) and unbatched condition cases.

Fixes #490; may also fix #1120.
vmap support / batching rule for lax.cond
It would be great to have vmap support for `lax.cond` similar to the recently added vmap support for `lax.while_loop`. I know that I can use `np.where` as a workaround for some cases, but for computationally expensive if conditions a batched cond might be better. 
Feature request: `vmap`-able boolean function options
I would like to write functions that take static options (associated with branching within the function) that are compatible with both `vmap` and `jit`. Some examples of things that currently work/don't work:

```python
def python_control_flow(x, option=False):
  return x**2 if option else x**3

vmap(python_control_flow)(np.arange(3))  # works
vmap(jit(python_control_flow))(np.arange(3))  # works
(lambda x: jit(python_control_flow)(x, True))(np.arange(3))  # doesn't work
vmap(jit(lambda x: python_control_flow(x, True)))(np.arange(3))  # works
vmap(python_control_flow, in_axes=(0, None))(np.arange(3), True)  # doesn't work

def cond_control_flow(x, option=False):
  return lax.cond(option, x, lambda x: x**2, x, lambda x: x**3)

cond_control_flow(np.arange(3))  # works
jit(cond_control_flow)(np.arange(3))  # doesn't work, https://github.com/google/jax/issues/1077
vmap(cond_control_flow)(np.arange(3))  # doesn't work
vmap(jit(cond_control_flow))(np.arange(3))  # doesn't work
vmap(cond_control_flow, in_axes=(0, None))(np.arange(3), True)  # doesn't work
```

It seems that for the best compatibility with `jit` the `lax.cond` control flow is preferable, perhaps with a wrapper, e.g.,
```python
@jit
def _cond_control_flow(x, option):
  return lax.cond(option, x, lambda x: x**2, x, lambda x: x**3)

def cond_control_flow(x, option=False):
  return _cond_control_flow(x, option)
```
to avoid https://github.com/google/jax/issues/1077, but no approach along these lines is currently compatible with a `vmap` down the line due to `NotImplementedError: Batching rule for 'cond' not implemented`.


Discussion:
What about batching via `lax.map`?
Apparently the right kind of lattice join is already being invoked in the right place, so the updated implementation both (a) works and (b) raises to the batch tracer only when necessary, even though the actual batch rule doesn't do any raising at all. (Earlier, I was cargo-culting off of the `while_loop` batching rule, which does need to do these things because there's no generic lattice join being called in `while_loop` itself).
Whoa! Do you mean the lattice join in [this code](https://github.com/google/jax/blob/a6ec5a916cae5b71e0232633b89e7490f77d81ef/jax/interpreters/partial_eval.py#L317) or elsewhere? How did you figure out that it wasn't doing any extra raising?
As I usually only need the if condition (no else part) I am considering a single iteration `while_loop` instead of `cond` as another workaround. What do you think?
I think this feature sounds super fun to implement and I can't wait to try it!

But in the meantime a single-iteration `while_loop` (or `fori_loop if it's more convenient) is probably a good workaround.


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
@@ -164,7 +164,7 @@ def while_loop(cond_fun, body_fun, init_val):
     raise TypeError(msg.format(cond_tree))
   if cond_jaxpr.out_avals != [ShapedArray((), onp.bool_)]:
     msg = "cond_fun must return a boolean scalar, but got output type(s) {}."
-    raise TypeError(msg.format(coud_jaxpr.out_avals))
+    raise TypeError(msg.format(cond_jaxpr.out_avals))
   if not treedef_children(in_tree) == [body_tree]:
     msg = "body_fun output pytree structure must match init_val, got {} and {}."
     raise TypeError(msg.format(body_tree, treedef_children(in_tree)[0]))
@@ -301,18 +301,6 @@ def cond(pred, true_operand, true_fun, false_operand, false_fun):
       true_nconsts=len(true_consts), false_nconsts=len(false_consts))
   return tree_unflatten(out_tree, out)
 
-def _cond_impl(pred, *args, **kwargs):
-  true_jaxpr, false_jaxpr, true_nconsts, false_nconsts = split_dict(
-      kwargs, ["true_jaxpr", "false_jaxpr", "true_nconsts", "false_nconsts"])
-  true_nops = len(true_jaxpr.in_avals) - true_nconsts
-  true_consts, true_ops, false_consts, false_ops = split_list(
-      args, [true_nconsts, true_nops, false_nconsts])
-
-  if pred:
-    return core.jaxpr_as_fun(true_jaxpr)(*(true_consts + true_ops))
-  else:
-    return core.jaxpr_as_fun(false_jaxpr)(*(false_consts + false_ops))
-
 def _cond_abstract_eval(*args, **kwargs):
   return kwargs["true_jaxpr"].out_avals
 
@@ -339,10 +327,47 @@ def make_computation(name, jaxpr, op_shape):
 
   return c.Conditional(pred, true_op, true_c, false_op, false_c)
 
+def _cond_batching_rule(args, dims, true_jaxpr, false_jaxpr, true_nconsts,
+                        false_nconsts):
+  # TODO: maybe avoid moving arg axes to front if we're promoting to select?
+  args = [batching.moveaxis(x, d, 0) if d is not batching.not_mapped and d != 0
+          else x for x, d in zip(args, dims)]
+  true_nops = len(true_jaxpr.in_avals) - true_nconsts
+  (pred,), true_consts, true_ops, false_consts, false_ops = split_list(
+      args, [1, true_nconsts, true_nops, false_nconsts])
+  size, = {x.shape[d] for x, d in zip(args, dims) if d is not batching.not_mapped}
+  orig_bat = [d is not batching.not_mapped for d in dims]
+  (pred_bat,), t_bat, tconst_bat, f_bat, fconst_bat = split_list(
+    orig_bat, [1, true_nconsts, len(true_ops), false_nconsts])
+
+  _, true_out_bat = batching.batch_jaxpr(true_jaxpr, size, tconst_bat + t_bat, False)
+  _, false_out_bat = batching.batch_jaxpr(false_jaxpr, size, fconst_bat + f_bat, False)
+  out_bat = [a or b for a, b in zip(true_out_bat, false_out_bat)]
+
+  true_jaxpr_batched, _ = batching.batch_jaxpr(true_jaxpr, size, tconst_bat + t_bat, out_bat)
+  false_jaxpr_batched, _ = batching.batch_jaxpr(false_jaxpr, size, fconst_bat + f_bat, out_bat)
+
+  if pred_bat:
+    true_out = core.jaxpr_as_fun(true_jaxpr_batched)(*(true_consts + true_ops))
+    false_out = core.jaxpr_as_fun(false_jaxpr_batched)(*(false_consts + false_ops))
+    true_out = [batching.broadcast(x, size, 0) if not b else x
+                for x, b in zip(true_out, out_bat)]
+    false_out = [batching.broadcast(x, size, 0) if not b else x
+                for x, b in zip(false_out, out_bat)]
+    return [lax.select(pred, t, f)
+            for t, f in zip(true_out, false_out)], [0] * len(true_out)
+  else:
+    out_dims = [0 if b else batching.not_mapped for b in out_bat]
+    return cond_p.bind(
+      *itertools.chain([pred], true_consts, true_ops, false_consts, false_ops),
+      true_jaxpr=true_jaxpr_batched, false_jaxpr=false_jaxpr_batched,
+      true_nconsts=len(true_consts), false_nconsts=len(false_consts)), out_dims
+
 cond_p = lax.Primitive('cond')
 cond_p.multiple_results = True
-cond_p.def_impl(_cond_impl)
+cond_p.def_impl(partial(xla.apply_primitive, cond_p))
 cond_p.def_abstract_eval(_cond_abstract_eval)
+batching.primitive_batchers[cond_p] = _cond_batching_rule
 xla.initial_style_translations[cond_p] = _cond_translation_rule
 
 


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
@@ -475,6 +475,58 @@ def cfun(x):
     self.assertEqual(fun(4), cfun(4))
     self.assertEqual(cfun(4), (4, 2., 4.))
 
+  def testCondBatched(self):
+    def fun(x, y, z):
+      pred = lax.lt(x, 3)
+      true_fun = lambda y: y
+      false_fun = lambda z: lax.neg(z)
+      return lax.cond(pred, y, true_fun, z, false_fun)
+
+    # these cases stay as cond
+    x = onp.array(2)
+    y = onp.array([1, 2])
+    z = onp.array([3, 4])
+    ans = api.vmap(fun, (None, 0, 0))(x, y, z)
+    jaxpr = api.make_jaxpr(api.vmap(fun, (None, 0, 0)))(x, y, z)
+    expected = onp.array([1, 2])
+    self.assertAllClose(ans, expected, check_dtypes=False)
+    assert "select" not in str(jaxpr)
+
+    x = onp.array(4)
+    ans = api.vmap(fun, (None, 0, 0))(x, y, z)
+    jaxpr = api.make_jaxpr(api.vmap(fun, (None, 0, 0)))(x, y, z)
+    expected = onp.array([-3, -4])
+    self.assertAllClose(ans, expected, check_dtypes=False)
+    assert "select" not in str(jaxpr)
+
+    fun = api.jit(fun)
+    ans = api.vmap(fun, (None, 0, 0))(x, y, z)
+    expected = onp.array([-3, -4])
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+    z = onp.array(5)
+    ans = api.vmap(fun, (None, 0, None))(x, y, z)
+    jaxpr = api.make_jaxpr(api.vmap(fun, (None, 0, None)))(x, y, z)
+    expected = onp.array([-5, -5])
+    self.assertAllClose(ans, expected, check_dtypes=False)
+    assert "select" not in str(jaxpr)
+
+
+    # these cases become select
+    x = onp.array([2, 4])
+    ans = api.vmap(fun, (0, 0, None))(x, y, z)
+    jaxpr = api.make_jaxpr(api.vmap(fun, (0, 0, None)))(x, y, z)
+    expected = onp.array([1, -5])
+    self.assertAllClose(ans, expected, check_dtypes=False)
+    assert "select" in str(jaxpr)
+
+    z = onp.array([3, 4])
+    ans = api.vmap(fun)(x, y, z)
+    jaxpr = api.make_jaxpr(api.vmap(fun))(x, y, z)
+    expected = onp.array([1, -4])
+    self.assertAllClose(ans, expected, check_dtypes=False)
+    assert "select" in str(jaxpr)
+
   def testIssue514(self):
     # just check this doesn't crash
     lax.cond(True,