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


Discussion:
Thanks for catching this! I think it's just that the `Select` HLO / `lax.select` primitive has limited automatic broadcasting (which we sometimes forget), so we need to do the broadcast ourselves. We did this [for the While batching rule](https://github.com/google/jax/blob/fd9e0333a6430bb06bd7a0c9faa650c66fdf2a21/jax/lax/lax_control_flow.py#L230-L237) but neglected to in the cond batching rule and [only called the raw select function without the broadcasting we need](https://github.com/google/jax/blob/fd9e0333a6430bb06bd7a0c9faa650c66fdf2a21/jax/lax/lax_control_flow.py#L359).
Thanks @mattjj! Following your suggestion, I am able to resolve it with
```
def _pred_bcast_select1(pred, x, y):
  import jax.numpy as np
  pred = np.broadcast_to(np.reshape(pred, np.shape(pred) + (1,) * (np.ndim(x) - np.ndim(pred))), np.shape(x))
  return lax.select(pred, x, y)
```
There is still a small issue of batching `cond` which throws
```
TypeError: broadcast_in_dim broadcast_dimensions must have length equal to operand ndim, got broadcast_dimensions (1, 2) for operand ndim 3.
```
at
```
~/jax/jax/lax/lax_control_flow.py in _cond_batching_rule(args, dims, true_jaxpr, false_jaxpr, true_nconsts, false_nconsts)
    356 
    357   if pred_bat:
--> 358     true_out = core.jaxpr_as_fun(true_jaxpr_batched)(*(true_consts + true_ops))
    359     false_out = core.jaxpr_as_fun(false_jaxpr_batched)(*(false_consts + false_ops))
    360     true_out = [batching.broadcast(x, size, 0) if not b else x
```
but it is harder to make a repro code. >"< I am working on it btw.
Here is another repro script which triggers another error but related to cond batching rule
```
import os; os.environ['XLA_FLAGS'] = '--xla_force_host_platform_device_count=2'
import jax
from jax.config import config; config.update('jax_platform_name', 'cpu')
import jax.numpy as np
from jax import random, lax, jit, vmap, pmap


def welford_covariance():
    def init_fn(size):
        return np.zeros(size), np.zeros(size), 0

    def update_fn(sample, state):
        mean, m2, n = state
        n = n + 1
        delta_pre = sample - mean
        mean = mean + delta_pre / n
        delta_post = sample - mean
        m2 = m2 + delta_pre * delta_post
        return mean, m2, n

    def final_fn(state):
        mean, m2, n = state
        cov = m2 / (n - 1)
        cov_inv_sqrt = np.sqrt(np.reciprocal(cov))
        return cov, cov_inv_sqrt

    return init_fn, update_fn, final_fn


def warmup_adapter():
    mm_init, mm_update, mm_final = welford_covariance()

    def init_fn(z, rng, mass_matrix_size):
        inverse_mass_matrix = np.ones(mass_matrix_size)
        mass_matrix_sqrt = inverse_mass_matrix
        mm_state = mm_init(mass_matrix_size)
        return (inverse_mass_matrix, mass_matrix_sqrt, mm_state, rng)

    def _update_at_window_end(z, rng_ss, state):
        inverse_mass_matrix, mass_matrix_sqrt, mm_state, rng = state
        inverse_mass_matrix, mass_matrix_sqrt = mm_final(mm_state)
        mm_state = mm_init(inverse_mass_matrix.shape[-1])
        return (inverse_mass_matrix, mass_matrix_sqrt, mm_state, rng)

    def update_fn(t, accept_prob, z, state):
        inverse_mass_matrix, mass_matrix_sqrt, mm_state, rng = state
        rng, rng_ss = random.split(rng)
        state = (inverse_mass_matrix, mass_matrix_sqrt, mm_state, rng)
        state = lax.cond(t < 10,
                         (z, rng_ss, state), lambda args: _update_at_window_end(*args),
                         state, lambda x: x)
        return state

    return init_fn, update_fn

def f(z, rng):
    wa_init, wa_update = warmup_adapter()
    wa_state = wa_init(z, rng, mass_matrix_size=3)
    wa_state = wa_update(1, 0.1, z, wa_state)
    return wa_state

pmap(f)(np.ones((2, 3)), random.split(random.PRNGKey(0)))  # pass
vmap(f)(np.ones((2, 3)), random.split(random.PRNGKey(0)))  # fail
```
Error details
<details>

```
---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
<ipython-input-3-3ada91d9ee7b> in <module>
----> 1 vmap(f)(np.ones((2, 3)), random.split(random.PRNGKey(0)))

~/jax/jax/api.py in batched_fun(*args)
    607     flat_fun, out_tree = flatten_fun_nokwargs(f, in_tree)
    608     out_flat = batching.batch(flat_fun, args_flat, _flatten_axes(in_tree, in_axes),
--> 609                               lambda: _flatten_axes(out_tree(), out_axes))
    610     return tree_unflatten(out_tree(), out_flat)
    611 

~/jax/jax/interpreters/batching.py in batch(fun, in_vals, in_dims, out_dim_dests)
     39 
     40 def batch(fun, in_vals, in_dims, out_dim_dests):
---> 41   out_vals, out_dims = batch_fun(fun, in_vals, in_dims)
     42   size, = {x.shape[d] for x, d in zip(in_vals, in_dims) if d is not not_mapped}
     43   return map(partial(matchaxis, size), out_dims, out_dim_dests(), out_vals)

~/jax/jax/interpreters/batching.py in batch_fun(fun, in_vals, in_dims)
     46   with new_master(BatchTrace) as master:
     47     fun, out_dims = batch_subtrace(fun, master, in_dims)
---> 48     out_vals = fun.call_wrapped(*in_vals)
     49     del master
     50   return out_vals, out_dims()

~/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:

<ipython-input-1-f5b4fcd16d7b> in f(z, rng)
     57     wa_init, wa_update = warmup_adapter()
     58     wa_state = wa_init(z, rng, mass_matrix_size=3)
---> 59     wa_state = wa_update(1, 0.1, z, wa_state)
     60     return wa_state

<ipython-input-1-f5b4fcd16d7b> in update_fn(t, accept_prob, z, state)
     49         state = lax.cond(t < 10,
     50                          (z, rng_ss, state), lambda args: _update_at_window_end(*args),
---> 51                          state, lambda x: x)
     52         return state
     53 

~/jax/jax/lax/lax_control_flow.py in cond(pred, true_operand, true_fun, false_operand, false_fun)
    300       *itertools.chain([pred], true_consts, true_ops, false_consts, false_ops),
    301       true_jaxpr=true_jaxpr, false_jaxpr=false_jaxpr,
--> 302       true_nconsts=len(true_consts), false_nconsts=len(false_consts))
    303   return tree_unflatten(out_tree, out)
    304 

~/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)
    369       *itertools.chain([pred], true_consts, true_ops, false_consts, false_ops),
    370       true_jaxpr=true_jaxpr_batched, false_jaxpr=false_jaxpr_batched,
--> 371       true_nconsts=len(true_consts), false_nconsts=len(false_consts)), out_dims
    372 
    373 cond_p = lax.Primitive('cond')

~/jax/jax/core.py in bind(self, *args, **kwargs)
    128     top_trace = find_top_trace(args)
    129     if top_trace is None:
--> 130       return self.impl(*args, **kwargs)
    131 
    132     tracers = map(top_trace.full_raise, args)

~/jax/jax/interpreters/xla.py in apply_primitive(prim, *args, **params)
    121   """Impl rule that compiles and runs a single primitive 'prim' using XLA."""
    122   abstract_args = map(abstractify, args)
--> 123   compiled_fun = xla_primitive_callable(prim, *abstract_args, **params)
    124   return compiled_fun(*args)
    125 

~/jax/jax/interpreters/xla.py in xla_primitive_callable(prim, *abstract_args, **params)
    134     handle_result = aval_to_result_handler(aval_out)
    135   xla_shapes = tuple(map(aval_to_xla_shape, abstract_args))
--> 136   built_c = primitive_computation(prim, *xla_shapes, **params)
    137   compiled = built_c.Compile(xla_shapes, xb.get_compile_options(),
    138                              backend=xb.get_backend(backend))

~/jax/jax/interpreters/xla.py in primitive_computation(prim, *xla_shapes, **params)
    157   elif prim in initial_style_translations:
    158     rule = initial_style_translations[prim]
--> 159     rule(c, AxisEnv(), *xla_args,  backend=backend, **new_params)  # side-effect on c
    160   else:
    161     raise NotImplementedError("XLA translation rule for {} not found".format(prim))

~/jax/jax/lax/lax_control_flow.py in _cond_translation_rule(c, axis_env, pred, *args, **kwargs)
    323 
    324   true_op = c.Tuple(*(true_consts + true_ops))
--> 325   true_c = make_computation("true_comp", true_jaxpr, c.GetShape(true_op))
    326 
    327   false_op = c.Tuple(*(false_consts + false_ops))

~/jax/jax/lax/lax_control_flow.py in make_computation(name, jaxpr, op_shape)
    319     ops = [c.GetTupleElement(op, i) for i in range(len(jaxpr.in_avals))]
    320     out = c.Call(xla.jaxpr_computation(jaxpr.jaxpr, backend, axis_env, jaxpr.literals, (),
--> 321                                        *_map(c.GetShape, ops)), ops)
    322     return c.Build(out)
    323 

~/jax/jax/interpreters/xla.py in jaxpr_computation(jaxpr, backend, axis_env, const_vals, freevar_shapes, *arg_shapes)
    233 def jaxpr_computation(jaxpr, backend, axis_env, const_vals, freevar_shapes, *arg_shapes):
    234   c, out_nodes = _jaxpr_computation(jaxpr, backend, axis_env, const_vals, freevar_shapes,
--> 235                                     *arg_shapes)
    236   return c.Build(c.Tuple(*out_nodes))
    237 

~/jax/jax/interpreters/xla.py in _jaxpr_computation(jaxpr, backend, axis_env, const_vals, freevar_shapes, *arg_shapes)
    303       raise NotImplementedError(msg.format(eqn.primitive.name))
    304 
--> 305     c.GetShape(ans)  # force xla to do shape error checking
    306     out_nodes = xla_destructure(c, ans) if eqn.primitive.multiple_results else [ans]
    307     _map(write, eqn.outvars, out_nodes)

~/miniconda3/envs/pydata/lib/python3.6/site-packages/jaxlib/xla_client.py in GetShape(self, operand)
    711 
    712   def GetShape(self, operand):
--> 713     return self._builder.GetShape(operand)
    714 
    715   def SetOpMetadata(self, op_metadata):

RuntimeError: Invalid argument: Size of broadcast_dimensions has to match operand's rank; operand rank: 2, size of broadcast_dimensions 1.
```

</details>

I think it is related to https://github.com/google/jax/pull/1237 but here the issue happens in `batching` context.
The above repro script triggers error at `pred_bat=False` branch, a similar issue happens at `pred_bat=True` branch too
```
def f(t, z, rng):
    wa_init, wa_update = warmup_adapter()
    wa_state = wa_init(z, rng, mass_matrix_size=3)
    wa_state = wa_update(t, 0.1, z, wa_state)
    return wa_state

vmap(f)(np.array([1, 2]), np.ones((2, 3)), random.split(random.PRNGKey(0)))
```

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,6 +329,10 @@ def make_computation(name, jaxpr, op_shape):
 
   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]


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
@@ -30,6 +30,7 @@
 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 @@ def fun(x, y, z):
     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,