Repository: google/jax
Issue: cond batching rule shape mismatch

A simple script to repro
```
from jax import lax, numpy as np, vmap, jit, random
from jax.config import config; config.update('jax_platform_name', 'cpu')

def f(rng, x):
    cond = random.bernoulli(rng)
    y = np.abs(x) - 1.
    return lax.cond(cond, x, lambda x: x, y, lambda x: x)

def body_fn(i, state):
    rng, x = state
    key, subkey = random.split(rng)
    return key, f(subkey, x)

def g(rng, x):
    return lax.fori_loop(0, 10, body_fn, (rng, x))

vmap(g)(random.split(random.PRNGKey(0), 3), np.ones((3, 4)))
```
which triggers
```
TypeError: select pred must be scalar or have the same shape as on_true and on_false, got pred shape (3,) for on_true and on_false of shape (3, 4).
```
More details

<details>

```sh
~/jax/jax/lax/lax_control_flow.py in fori_loop(lower, upper, body_fun, init_val)
    116     return lax.add(i, lax._const(i, 1)), body_fun(i, x)
    117 
--> 118   _, result = while_loop(while_cond_fun, while_body_fun, (lower, init_val))
    119   return result
    120 

~/jax/jax/lax/lax_control_flow.py in while_loop(cond_fun, body_fun, init_val)
    171   outs = while_p.bind(*itertools.chain(cond_consts, body_consts, init_vals),
    172                       cond_nconsts=len(cond_consts), cond_jaxpr=cond_jaxpr,
--> 173                       body_nconsts=len(body_consts), body_jaxpr=body_jaxpr)
    174   return tree_unflatten(body_tree, outs)
    175 

~/jax/jax/core.py in bind(self, *args, **kwargs)
    131 
    132     tracers = map(top_trace.full_raise, args)
--> 133     out_tracer = top_trace.process_primitive(self, tracers, kwargs)
    134     if self.multiple_results:
    135       return map(full_lower, out_tracer)

~/jax/jax/interpreters/batching.py in process_primitive(self, primitive, tracers, params)
    113       # TODO(mattjj,phawkins): if no rule implemented, could vmap-via-map here
    114       batched_primitive = get_primitive_batcher(primitive)
--> 115       val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
    116       if primitive.multiple_results:
    117         return map(partial(BatchTracer, self), val_out, dim_out)

~/jax/jax/lax/lax_control_flow.py in _while_loop_batching_rule(args, dims, cond_nconsts, cond_jaxpr, body_nconsts, body_jaxpr)
    247     batched = bconst_bat + carry_bat
    248     body_jaxpr_batched, carry_bat_out = batching.batch_jaxpr(
--> 249         body_jaxpr, size, batched, instantiate=carry_bat)
    250     cond_jaxpr_batched, (pred_bat,) = batching.batch_jaxpr(
    251         cond_jaxpr, size, cconst_bat + carry_bat, instantiate=False)

~/jax/jax/interpreters/batching.py in batch_jaxpr(jaxpr, size, batched, instantiate)
    293               for a, b in zip(jaxpr.in_avals, batched)]
    294   in_pvals = [pe.PartialVal((aval, core.unit)) for aval in avals_in]
--> 295   jaxpr_out, pvals_out, consts_out = pe.trace_to_jaxpr(f, in_pvals, instantiate=True)
    296   avals_out, _ = unzip2(pvals_out)
    297   jaxpr_out = core.TypedJaxpr(jaxpr_out, consts_out, avals_in, avals_out)

~/jax/jax/interpreters/partial_eval.py in trace_to_jaxpr(fun, pvals, **kwargs)
    316   with new_master(JaxprTrace) as master:
    317     fun = trace_to_subjaxpr(fun, master, instantiate)
--> 318     jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
    319     assert not env
    320     del master

~/jax/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    163 
    164     del gen
--> 165     ans = self.f(*args, **dict(self.params, **kwargs))
    166     del args
    167     while stack:

~/jax/jax/core.py in jaxpr_as_fun(typed_jaxpr, *args)
     77 @curry
     78 def jaxpr_as_fun(typed_jaxpr, *args):
---> 79   return eval_jaxpr(typed_jaxpr.jaxpr, typed_jaxpr.literals, (), *args)
     80 
     81 

~/jax/jax/core.py in eval_jaxpr(jaxpr, consts, freevar_vals, *args)
    183                in eqn.bound_subjaxprs]
    184     subfuns = map(lu.wrap_init, subfuns)
--> 185     ans = eqn.primitive.bind(*(subfuns + in_vals), **eqn.params)
    186     if eqn.primitive.multiple_results:
    187       map(write, eqn.outvars, ans)

~/jax/jax/core.py in bind(self, *args, **kwargs)
    131 
    132     tracers = map(top_trace.full_raise, args)
--> 133     out_tracer = top_trace.process_primitive(self, tracers, kwargs)
    134     if self.multiple_results:
    135       return map(full_lower, out_tracer)

~/jax/jax/interpreters/batching.py in process_primitive(self, primitive, tracers, params)
    113       # TODO(mattjj,phawkins): if no rule implemented, could vmap-via-map here
    114       batched_primitive = get_primitive_batcher(primitive)
--> 115       val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
    116       if primitive.multiple_results:
    117         return map(partial(BatchTracer, self), val_out, dim_out)

~/jax/jax/lax/lax_control_flow.py in _cond_batching_rule(args, dims, true_jaxpr, false_jaxpr, true_nconsts, false_nconsts)
    358                 for x, b in zip(false_out, out_bat)]
    359     return [lax.select(pred, t, f)
--> 360             for t, f in zip(true_out, false_out)], [0] * len(true_out)
    361   else:
    362     out_dims = [0 if b else batching.not_mapped for b in out_bat]

~/jax/jax/lax/lax_control_flow.py in <listcomp>(.0)
    358                 for x, b in zip(false_out, out_bat)]
    359     return [lax.select(pred, t, f)
--> 360             for t, f in zip(true_out, false_out)], [0] * len(true_out)
    361   else:
    362     out_dims = [0 if b else batching.not_mapped for b in out_bat]

~/jax/jax/lax/lax.py in select(pred, on_true, on_false)
    611   operator.
    612   """
--> 613   return select_p.bind(pred, on_true, on_false)
    614 
    615 def slice(operand, start_indices, limit_indices, strides=None):

~/jax/jax/core.py in bind(self, *args, **kwargs)
    131 
    132     tracers = map(top_trace.full_raise, args)
--> 133     out_tracer = top_trace.process_primitive(self, tracers, kwargs)
    134     if self.multiple_results:
    135       return map(full_lower, out_tracer)

~/jax/jax/interpreters/partial_eval.py in process_primitive(self, primitive, tracers, params)
     93       tracers = map(self.instantiate_const, tracers)
     94       avals = [t.aval for t in tracers]
---> 95       out_aval = primitive.abstract_eval(*avals, **params)
     96       # TODO(dougalm): think about whether these ref cycles will leak memory
     97       if primitive.multiple_results:

~/jax/jax/lax/lax.py in standard_abstract_eval(shape_rule, dtype_rule, *args, **kwargs)
   1468     return ShapedArray(shape_rule(*args, **kwargs), dtype_rule(*args, **kwargs))
   1469   elif least_specialized is ShapedArray:
-> 1470     return ShapedArray(shape_rule(*args, **kwargs), dtype_rule(*args, **kwargs))
   1471   elif least_specialized is UnshapedArray:
   1472     return UnshapedArray(dtype_rule(*args, **kwargs))

~/jax/jax/lax/lax.py in _select_shape_rule(pred, on_true, on_false)
   2561     msg = ("select pred must be scalar or have the same shape as on_true and "
   2562            "on_false, got pred shape {} for on_true and on_false of shape {}.")
-> 2563     raise TypeError(msg.format(pred.shape, on_true.shape))
   2564   return on_true.shape
   2565 
```
  </details>

I think that the logic of `cond_batching_rule` is correct but the `pred` shape is somehow not broadcasted in `_cond_batching_rule`.
cc @jekbradbury @mattjj @neerajprad

Changed files:
jax/lax/lax_control_flow.py, tests/lax_control_flow_test.py

diff --git a/jax/lax/lax_control_flow.py b/jax/lax/lax_control_flow.py
index d9689500a..dbd1d09e0 100644
--- a/jax/lax/lax_control_flow.py
+++ b/jax/lax/lax_control_flow.py
@@ -329,6 +329,10 @@ def _cond_translation_rule(c, axis_env, pred, *args, **kwargs):
 
   return c.Conditional(pred, true_op, true_c, false_op, false_c)
 
+def _cond_pred_bcast_select(pred, x, y):
+  bcast_pred = lax.broadcast_in_dim(pred, onp.shape(x), list(range(onp.ndim(pred))))
+  return lax.select(bcast_pred, x, y)
+
 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?
@@ -339,8 +343,8 @@ def _cond_batching_rule(args, dims, true_jaxpr, false_jaxpr, true_nconsts,
       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])
+  (pred_bat,), tconst_bat, t_bat, fconst_bat, f_bat = split_list(
+    orig_bat, [1, true_nconsts, true_nops, 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)
@@ -355,8 +359,8 @@ def _cond_batching_rule(args, dims, true_jaxpr, false_jaxpr, true_nconsts,
     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 x, b in zip(false_out, out_bat)]
+    return [_cond_pred_bcast_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]
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
index b2f3f13fc..fc5ef8aa2 100644
--- a/tests/lax_control_flow_test.py
+++ b/tests/lax_control_flow_test.py
@@ -30,6 +30,7 @@ import numpy.random as npr
 from jax import api
 from jax import core
 from jax import lax
+from jax import random
 from jax import test_util as jtu
 from jax.util import unzip2
 from jax.lib import xla_bridge
@@ -527,6 +528,21 @@ class LaxControlFlowTest(jtu.JaxTestCase):
     self.assertAllClose(ans, expected, check_dtypes=False)
     assert "select" in str(jaxpr)
 
+  def testIssue1263(self):
+    def f(rng, x):
+      cond = random.bernoulli(rng)
+      return lax.cond(cond, x, lambda x: x, np.abs(x) - 1., lambda x: x)
+
+    def body_fn(i, state):
+      rng, x = state
+      key, subkey = random.split(rng)
+      return key, f(subkey, x)
+
+    def g(rng, x):
+      return lax.fori_loop(0, 10, body_fn, (rng, x))
+
+    api.vmap(g)(random.split(random.PRNGKey(0), 3), np.ones((3, 4)))
+
   def testIssue514(self):
     # just check this doesn't crash
     lax.cond(True,