Repository: google/jax

Issue:
Using `vmap` on `grad` on function with `lax.cond` throws errors
Hi, 

I'm using `vmap` on `grad` on functions containing `lax.cond`, this trows errors. Is it supposed to work yet? If I switch to the equivalent `np.where` there's no problem:

```python
def f_1(x):
    return x ** 2

def f_2(x):
    return x ** 3


def f(x):
    y = lax.cond(x > 0, x, f_1, x, f_2)
    return y

def g(x):
    y = np.where(x > 0, f_1(x), f_2(x))
    return y

x = np.linspace(-1, 1, 20)
y = vmap(grad(f))(x)
print(y)
```


<details>
---------------------------------------------------------------------------
IndexError                                Traceback (most recent call last)
<ipython-input-55-e5aab592a69e> in <module>
     15 
     16 x = np.linspace(-1, 1, 20)
---> 17 y = vmap(grad(f))(x)
     18 print(y)

~/one-leg/.venv/lib/python3.5/site-packages/jax/api.py in batched_fun(*args)
    690     _check_axis_sizes(in_tree, args_flat, in_axes_flat)
    691     out_flat = batching.batch(flat_fun, args_flat, in_axes_flat,
--> 692                               lambda: _flatten_axes(out_tree(), out_axes))
    693     return tree_unflatten(out_tree(), out_flat)
    694 

~/one-leg/.venv/lib/python3.5/site-packages/jax/interpreters/batching.py in batch(fun, in_vals, in_dims, out_dim_dests)
     35 def batch(fun, in_vals, in_dims, out_dim_dests):
     36   size, = {x.shape[d] for x, d in zip(in_vals, in_dims) if d is not not_mapped}
---> 37   out_vals, out_dims = batch_fun(fun, in_vals, in_dims)
     38   return map(partial(matchaxis, size), out_dims, out_dim_dests(), out_vals)
     39 

~/one-leg/.venv/lib/python3.5/site-packages/jax/interpreters/batching.py in batch_fun(fun, in_vals, in_dims)
     41   with new_master(BatchTrace) as master:
     42     fun, out_dims = batch_subtrace(fun, master, in_dims)
---> 43     out_vals = fun.call_wrapped(*in_vals)
     44     del master
     45   return out_vals, out_dims()

~/one-leg/.venv/lib/python3.5/site-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    147     gen = None
    148 
--> 149     ans = self.f(*args, **dict(self.params, **kwargs))
    150     del args
    151     while stack:

~/one-leg/.venv/lib/python3.5/site-packages/jax/api.py in grad_f(*args, **kwargs)
    350   @wraps(fun, docstr=docstr, argnums=argnums)
    351   def grad_f(*args, **kwargs):
--> 352     _, g = value_and_grad_f(*args, **kwargs)
    353     return g
    354 

~/one-leg/.venv/lib/python3.5/site-packages/jax/api.py in value_and_grad_f(*args, **kwargs)
    405     f_partial, dyn_args = _argnums_partial(f, argnums, args)
    406     if not has_aux:
--> 407       ans, vjp_py = vjp(f_partial, *dyn_args)
    408     else:
    409       ans, vjp_py, aux = vjp(f_partial, *dyn_args, has_aux=True)

~/one-leg/.venv/lib/python3.5/site-packages/jax/api.py in vjp(fun, *primals, **kwargs)
   1283   if not has_aux:
   1284     flat_fun, out_tree = flatten_fun_nokwargs(fun, in_tree)
-> 1285     out_primal, out_vjp = ad.vjp(flat_fun, primals_flat)
   1286     out_tree = out_tree()
   1287   else:

~/one-leg/.venv/lib/python3.5/site-packages/jax/interpreters/ad.py in vjp(traceable, primals, has_aux)
    104 def vjp(traceable, primals, has_aux=False):
    105   if not has_aux:
--> 106     out_primals, pvals, jaxpr, consts = linearize(traceable, *primals)
    107   else:
    108     out_primals, pvals, jaxpr, consts, aux = linearize(traceable, *primals, has_aux=True)

~/one-leg/.venv/lib/python3.5/site-packages/jax/interpreters/ad.py in linearize(traceable, *primals, **kwargs)
     93   _, in_tree = tree_flatten(((primals, primals), {}))
     94   jvpfun_flat, out_tree = flatten_fun(jvpfun, in_tree)
---> 95   jaxpr, out_pvals, consts = pe.trace_to_jaxpr(jvpfun_flat, in_pvals)
     96   pval_primals, pval_tangents = tree_unflatten(out_tree(), out_pvals)
     97   aval_primals, const_primals = unzip2(pval_primals)

~/one-leg/.venv/lib/python3.5/site-packages/jax/interpreters/partial_eval.py in trace_to_jaxpr(fun, pvals, instantiate, stage_out_calls)
    352   with new_master(trace_type) as master:
    353     fun = trace_to_subjaxpr(fun, master, instantiate)
--> 354     jaxpr, (out_pvals, consts, env) = fun.call_wrapped(pvals)
    355     assert not env
    356     del master

~/one-leg/.venv/lib/python3.5/site-packages/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
    147     gen = None
    148 
--> 149     ans = self.f(*args, **dict(self.params, **kwargs))
    150     del args
    151     while stack:

<ipython-input-55-e5aab592a69e> in f(x)
      7 
      8 def f(x):
----> 9     y = lax.cond(x > 0, x, f_1, x, f_2)
     10     return y
     11 

~/one-leg/.venv/lib/python3.5/site-packages/jax/lax/lax_control_flow.py in cond(pred, true_operand, true_fun, false_operand, false_fun)
    421   out = cond_p.bind(
    422       *itertools.chain([pred], true_consts, true_ops, false_consts, false_ops),
--> 423       true_jaxpr=true_jaxpr, false_jaxpr=false_jaxpr, linear=linear)
    424   return tree_unflatten(true_out_tree, out)
    425 

~/one-leg/.venv/lib/python3.5/site-packages/jax/lax/lax_control_flow.py in cond_bind(true_jaxpr, false_jaxpr, linear, *args)
    675     core.check_jaxpr(false_jaxpr.jaxpr)
    676   return core.Primitive.bind(cond_p, *args, true_jaxpr=true_jaxpr,
--> 677                              false_jaxpr=false_jaxpr, linear=linear)
    678 
    679 cond_p = lax.Primitive('cond')

~/one-leg/.venv/lib/python3.5/site-packages/jax/core.py in bind(self, *args, **kwargs)
    180 
    181     tracers = map(top_trace.full_raise, args)
--> 182     out_tracer = top_trace.process_primitive(self, tracers, kwargs)
    183     if self.multiple_results:
    184       return map(full_lower, out_tracer)

~/one-leg/.venv/lib/python3.5/site-packages/jax/interpreters/ad.py in process_primitive(self, primitive, tracers, params)
    303           "Forward-mode differentiation rule for '{}' not implemented"
    304           .format(primitive))
--> 305     primal_out, tangent_out = jvp(primals_in, tangents_in, **params)
    306     if primitive.multiple_results:
    307       return [JVPTracer(self, x, t) for x, t in zip(primal_out, tangent_out)]

~/one-leg/.venv/lib/python3.5/site-packages/jax/lax/lax_control_flow.py in _cond_jvp(primals, tangents, true_jaxpr, false_jaxpr, linear)
    510   out = cond_p.bind(
    511       *itertools.chain([pred], tops, tops_dot, fops, fops_dot),
--> 512       true_jaxpr=true_jvp, false_jaxpr=false_jvp, linear=linear_jvp)
    513   out_primals, out_tangents = split_list(out, [len(out_nz)])
    514   out_tangents_iter = iter(out_tangents)

~/one-leg/.venv/lib/python3.5/site-packages/jax/lax/lax_control_flow.py in cond_bind(true_jaxpr, false_jaxpr, linear, *args)
    675     core.check_jaxpr(false_jaxpr.jaxpr)
    676   return core.Primitive.bind(cond_p, *args, true_jaxpr=true_jaxpr,
--> 677                              false_jaxpr=false_jaxpr, linear=linear)
    678 
    679 cond_p = lax.Primitive('cond')

~/one-leg/.venv/lib/python3.5/site-packages/jax/core.py in bind(self, *args, **kwargs)
    180 
    181     tracers = map(top_trace.full_raise, args)
--> 182     out_tracer = top_trace.process_primitive(self, tracers, kwargs)
    183     if self.multiple_results:
    184       return map(full_lower, out_tracer)

~/one-leg/.venv/lib/python3.5/site-packages/jax/interpreters/partial_eval.py in process_primitive(self, primitive, tracers, params)
     94   def process_primitive(self, primitive, tracers, params):
     95     if primitive in custom_partial_eval_rules:
---> 96       return custom_partial_eval_rules[primitive](self, *tracers, **params)
     97     else:
     98       return self.default_process_primitive(primitive, tracers, params)

~/one-leg/.venv/lib/python3.5/site-packages/jax/lax/lax_control_flow.py in _cond_partial_eval(trace, true_jaxpr, false_jaxpr, linear, *tracers)
    564   out_consts_res = cond_p.bind(
    565       *in_consts, true_jaxpr=true_jaxpr_1, false_jaxpr=false_jaxpr_1,
--> 566       linear=linear)
    567   out_consts, res = split_list(out_consts_res, [len(out_consts_res) - num_res])
    568 

~/one-leg/.venv/lib/python3.5/site-packages/jax/lax/lax_control_flow.py in cond_bind(true_jaxpr, false_jaxpr, linear, *args)
    675     core.check_jaxpr(false_jaxpr.jaxpr)
    676   return core.Primitive.bind(cond_p, *args, true_jaxpr=true_jaxpr,
--> 677                              false_jaxpr=false_jaxpr, linear=linear)
    678 
    679 cond_p = lax.Primitive('cond')

~/one-leg/.venv/lib/python3.5/site-packages/jax/core.py in bind(self, *args, **kwargs)
    180 
    181     tracers = map(top_trace.full_raise, args)
--> 182     out_tracer = top_trace.process_primitive(self, tracers, kwargs)
    183     if self.multiple_results:
    184       return map(full_lower, out_tracer)

~/one-leg/.venv/lib/python3.5/site-packages/jax/interpreters/batching.py in process_primitive(self, primitive, tracers, params)
    109       # TODO(mattjj,phawkins): if no rule implemented, could vmap-via-map here
    110       batched_primitive = get_primitive_batcher(primitive)
--> 111       val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
    112       if primitive.multiple_results:
    113         return map(partial(BatchTracer, self), val_out, dim_out)

~/one-leg/.venv/lib/python3.5/site-packages/jax/lax/lax_control_flow.py in _cond_batching_rule(***failed resolving arguments***)
    478                  for x, b in zip(false_out, out_bat)]
    479     return [_cond_pred_bcast_select(pred, t, f)
--> 480             for t, f in zip(true_out, false_out)], [0] * len(true_out)
    481   else:
    482     out_dims = [0 if b else batching.not_mapped for b in out_bat]

~/one-leg/.venv/lib/python3.5/site-packages/jax/lax/lax_control_flow.py in <listcomp>(.0)
    478                  for x, b in zip(false_out, out_bat)]
    479     return [_cond_pred_bcast_select(pred, t, f)
--> 480             for t, f in zip(true_out, false_out)], [0] * len(true_out)
    481   else:
    482     out_dims = [0 if b else batching.not_mapped for b in out_bat]

~/one-leg/.venv/lib/python3.5/site-packages/jax/lax/lax_control_flow.py in _cond_pred_bcast_select(pred, x, y)
    450 
    451 def _cond_pred_bcast_select(pred, x, y):
--> 452   bcast_pred = lax.broadcast_in_dim(pred, onp.shape(x), list(range(onp.ndim(pred))))
    453   return lax.select(bcast_pred, x, y)
    454 

~/one-leg/.venv/lib/python3.5/site-packages/jax/lax/lax.py in broadcast_in_dim(operand, shape, broadcast_dimensions)
    621   return broadcast_in_dim_p.bind(
    622       operand, shape=tuple(shape),
--> 623       broadcast_dimensions=tuple(broadcast_dimensions))
    624 
    625 def reshape(operand, new_sizes, dimensions=None):

~/one-leg/.venv/lib/python3.5/site-packages/jax/core.py in bind(self, *args, **kwargs)
    177     top_trace = find_top_trace(args)
    178     if top_trace is None:
--> 179       return self.impl(*args, **kwargs)
    180 
    181     tracers = map(top_trace.full_raise, args)

~/one-leg/.venv/lib/python3.5/site-packages/jax/lax/lax.py in _broadcast_in_dim_impl(operand, shape, broadcast_dimensions)
   2342   if type(operand) is xla.DeviceArray:
   2343     aval = ShapedArray(shape, _dtype(operand))
-> 2344     lazy_expr = lazy.broadcast(operand._lazy_expr, shape, broadcast_dimensions)
   2345     return xla.DeviceArray(aval, None, lazy_expr, operand.device_buffer)
   2346   else:

~/one-leg/.venv/lib/python3.5/site-packages/jax/lazy.py in broadcast(lexpr, shape, broadcast_dimensions)
    127   new_dims = [None] * len(shape)
    128   for i, d in enumerate(broadcast_dimensions):
--> 129     new_dims[d] = lexpr.dims[i]
    130   return LazyExpr(lexpr.input, shape, tuple(new_dims))
    131 

IndexError: list assignment index out of range
</details>

In a more complex application I have this error: 
```ValueError: broadcast_in_dim broadcast dimensions must be less than ndim(shape), got (0,) for shape ().```
If you want, I can share that code as well, but it's a bit more complicated than my toy example above.

Thanks for your time,

Rembert


Discussion:
It's supposed to work, but grad-of-cond is only pretty lightly tested at this point. Looks like you've found a bug for us!

Toy examples are perfect, btw. Thanks!

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
@@ -449,8 +449,11 @@ 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)
+  if core.get_aval(x) is core.get_aval(y) is core.abstract_unit:
+    return x
+  else:
+    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, linear):
   # TODO: maybe avoid moving arg axes to front if we're promoting to select?


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
@@ -820,6 +820,19 @@ def f(x):
     expected = f(4.)
     self.assertAllClose(y, expected, check_dtypes=False)
 
+  def testCondVmapGrad(self):
+    # https://github.com/google/jax/issues/2264
+    def f_1(x): return x ** 2
+    def f_2(x): return x ** 3
+
+    def f(x): return lax.cond(x > 0, x, f_1, x, f_2)
+    def g(x): return np.where(x > 0, f_1(x), f_2(x))
+
+    x = np.linspace(-1, 1, 20)
+    ans = api.vmap(api.grad(f))(x)
+    expected = api.vmap(api.grad(g))(x)
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
   def testIssue1263(self):
     def f(rng, x):
       cond = random.bernoulli(rng)