Repository: google/jax

Issue:
expm shape tracing issue through two vmaps and jvp
Hi,

I've been trying to narrow down the issue that arises when a vmap is applied to jvp that goes through a calculation which has its own vmap.  Although I am new to Jax, and potentially, lacking understanding of shape tracing but it seems to me that the problem occurs due to expm matrix operation. 

Consider the following code that either implements a Jacobian calculation for the matrix product operation:

> e^{I*(w[0,0]+w[0,1]})e^{I*(w[1,0]+w[1,1]}) 

or

> (I * w[0,0]) * (I * w[0,1]) * (I * w[1,0]) * (I * w[1,1])

where I is 2x2 identity matrix and w is a set of weights stored as a 2x2 array. The executed calculation can be toggled by commenting/uncommenting 2 lines in matrix_operation(). The code works well for the 2nd operation but breaks down for the 1st one throwing an assertion error related to the shape detected in lax_control_flow.py.

I am wondering whether it is a bug in expm or I am missing something? I would appreciate any help!

```
import jax.numpy as jnp
from jax import vmap, jvp
from functools import partial
from jax.numpy.linalg import multi_dot

def matrix_operation(w):
    assert w.shape == (2,)
    A = jnp.identity(2)
    return la.expm(A*(w[0]+w[1])) # -----  breaks down with this line
    #return jnp.matmul(A*w[0], A*w[1])  #----- works with this line! 
    
def cost_function(ws):
    same_w = vmap(matrix_operation)(ws)
    product = multi_dot(same_w)
    return product

def pushfwd_(func, weights, tangent): 
    return jvp(func, (weights,), (tangent,))

r = 2
c = 2
weights = jnp.ones((r, c))

pushfwd = partial(pushfwd_, cost_function, weights)

# a set of vectors with a single non-zero entry equal to 1
# and same shape as weights
tangents = jnp.reshape(jnp.identity(r*c), (c*r, r, c))

print(vmap(pushfwd)(tangents)) # this breaks with la.expm above
#pushfwd(tangents[0]) # this works with la.expm above
```



Discussion:
Wondering if it is related to the  issue #3056 ?
What specific error message do you see?

The short answer for what's wrong here is probably that JAX's `expm` doesn't support reverse-mode differentiation yet.
Thanks for the reply!

> The short answer for what's wrong here is probably that JAX's expm doesn't support reverse-mode differentiation yet.

Note that the code doesn't use reverse-mode differentiation, it runs JVP instead which is implemented. For example, if the last print is replaced with pushfwd(tangents[0]), all is good.

Also the code runs if the two body lines in cost_function(ws) are replaced with an explicit loop computing the product. So it seems that it is really the combination of two vmaps and jvp, that is causing the issue to arise. 

> What specific error message do you see?

Here is the full message. 

```
---------------------------------------------------------------------------
> AssertionError                            Traceback (most recent call last)
> <ipython-input-37-192dd72c18c4> in <module>
>      23 tangents = jnp.reshape(jnp.identity(r*c), (c*r, r, c))
>      24 
> ---> 25 print(vmap(pushfwd)(tangents))
> 
> ~/Codes/JAX_source/jax/jax/api.py in batched_fun(*args)
>     765     in_axes_flat = _flatten_axes(in_tree, in_axes)
>     766     _ = _mapped_axis_size(in_tree, args_flat, in_axes_flat, "vmap")
> --> 767     out_flat = batching.batch(flat_fun, args_flat, in_axes_flat,
>     768                               lambda: _flatten_axes(out_tree(), out_axes))
>     769     return tree_unflatten(out_tree(), out_flat)
> 
> ~/Codes/JAX_source/jax/jax/interpreters/batching.py in batch(fun, in_vals, in_dims, out_dim_dests)
>      32   # executes a batched version of `fun` following out_dim_dests
>      33   batched_fun = batch_fun(fun, in_dims, out_dim_dests)
> ---> 34   return batched_fun.call_wrapped(*in_vals)
>      35 
>      36 @lu.transformation_with_aux
> 
> ~/Codes/JAX_source/jax/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
>     148     gen = None
>     149 
> --> 150     ans = self.f(*args, **dict(self.params, **kwargs))
>     151     del args
>     152     while stack:
> 
> <ipython-input-37-192dd72c18c4> in pushfwd_(func, weights, tangent)
>      11 
>      12 def pushfwd_(func, weights, tangent):
> ---> 13     return jvp(func, (weights,), (tangent,))
>      14 
>      15 r = 2
> 
> ~/Codes/JAX_source/jax/jax/api.py in jvp(fun, primals, tangents)
>    1275   """
>    1276   _check_callable(fun)
> -> 1277   return _jvp(lu.wrap_init(fun), primals, tangents)
>    1278 
>    1279 def _jvp(fun: lu.WrappedFun, primals, tangents):
> 
> ~/Codes/JAX_source/jax/jax/api.py in _jvp(fun, primals, tangents)
>    1298       raise TypeError(msg.format(_dtype(p), _dtype(t)))
>    1299   flat_fun, out_tree = flatten_fun_nokwargs(fun, tree_def)
> -> 1300   out_primals, out_tangents = ad.jvp(flat_fun).call_wrapped(ps_flat, ts_flat)
>    1301   return (tree_unflatten(out_tree(), out_primals),
>    1302           tree_unflatten(out_tree(), out_tangents))
> 
> ~/Codes/JAX_source/jax/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
>     148     gen = None
>     149 
> --> 150     ans = self.f(*args, **dict(self.params, **kwargs))
>     151     del args
>     152     while stack:
> 
> <ipython-input-37-192dd72c18c4> in cost_function(ws)
>       6 
>       7 def cost_function(ws):
> ----> 8     same_w = vmap(matrix_operation)(ws)
>       9     product = la.multi_dot(same_w)
>      10     return product
> 
> ~/Codes/JAX_source/jax/jax/api.py in batched_fun(*args)
>     765     in_axes_flat = _flatten_axes(in_tree, in_axes)
>     766     _ = _mapped_axis_size(in_tree, args_flat, in_axes_flat, "vmap")
> --> 767     out_flat = batching.batch(flat_fun, args_flat, in_axes_flat,
>     768                               lambda: _flatten_axes(out_tree(), out_axes))
>     769     return tree_unflatten(out_tree(), out_flat)
> 
> ~/Codes/JAX_source/jax/jax/interpreters/batching.py in batch(fun, in_vals, in_dims, out_dim_dests)
>      32   # executes a batched version of `fun` following out_dim_dests
>      33   batched_fun = batch_fun(fun, in_dims, out_dim_dests)
> ---> 34   return batched_fun.call_wrapped(*in_vals)
>      35 
>      36 @lu.transformation_with_aux
> 
> ~/Codes/JAX_source/jax/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
>     148     gen = None
>     149 
> --> 150     ans = self.f(*args, **dict(self.params, **kwargs))
>     151     del args
>     152     while stack:
> 
> <ipython-input-37-192dd72c18c4> in matrix_operation(w)
>       2     assert w.shape == (2,)
>       3     A = jnp.identity(2)
> ----> 4     return la.expm(A*(w[0]+w[1])) # -----  breaks down with this line
>       5     #return jnp.matmul(A*w[0], A*w[1])  #----- works with this line!
>       6 
> 
> ~/Codes/JAX_source/jax/jax/api.py in __call__(self, *args)
>    1727     with core.initial_style_staging():
>    1728       jaxpr, _, consts = pe.trace_to_jaxpr(flat_fun, in_pvals, instantiate=True)
> -> 1729     outs = self.prim.bind(*it.chain(consts, args_flat), jaxpr=jaxpr,
>    1730                           in_tree=in_tree, out_tree=out_tree(),
>    1731                           num_consts=len(consts))
> 
> ~/Codes/JAX_source/jax/jax/core.py in bind(self, *args, **kwargs)
>     212 
>     213     tracers = map(top_trace.full_raise, args)
> --> 214     out_tracer = top_trace.process_primitive(self, tracers, kwargs)
>     215     if self.multiple_results:
>     216       return map(full_lower, out_tracer)
> 
> ~/Codes/JAX_source/jax/jax/interpreters/batching.py in process_primitive(self, primitive, tracers, params)
>     132       # TODO(mattjj,phawkins): if no rule implemented, could vmap-via-map here
>     133       batched_primitive = get_primitive_batcher(primitive)
> --> 134       val_out, dim_out = batched_primitive(vals_in, dims_in, **params)
>     135       if primitive.multiple_results:
>     136         return map(partial(BatchTracer, self), val_out, dim_out)
> 
> ~/Codes/JAX_source/jax/jax/api.py in fun_batch(args, dims, **params)
>    1750   def fun_batch(args, dims, **params):
>    1751     batched, out_dims = batching.batch_fun2(lu.wrap_init(fun_impl, params), dims)
> -> 1752     return batched.call_wrapped(*args), out_dims()
>    1753   batching.primitive_batchers[fun_p] = fun_batch
>    1754 
> 
> ~/Codes/JAX_source/jax/jax/linear_util.py in call_wrapped(self, *args, **kwargs)
>     148     gen = None
>     149 
> --> 150     ans = self.f(*args, **dict(self.params, **kwargs))
>     151     del args
>     152     while stack:
> 
> ~/Codes/JAX_source/jax/jax/api.py in fun_impl(*args, **params)
>    1741   def fun_impl(*args, **params):
>    1742     consts, args = split_list(args, [params['num_consts']])
> -> 1743     return core.eval_jaxpr(params['jaxpr'], consts, *args)
>    1744   fun_p.def_impl(fun_impl)
>    1745 
> 
> ~/Codes/JAX_source/jax/jax/core.py in eval_jaxpr(jaxpr, consts, *args)
>     279     else:
>     280       subfuns = []
> --> 281     ans = eqn.primitive.bind(*(subfuns + in_vals), **params)
>     282     if eqn.primitive.multiple_results:
>     283       map(write, eqn.outvars, ans)
> 
> ~/Codes/JAX_source/jax/jax/core.py in _call_bind(processor, post_processor, primitive, f, *args, **params)
>    1018     tracers = map(top_trace.full_raise, args)
>    1019     process = getattr(top_trace, processor)
> -> 1020     outs = map(full_lower, process(primitive, f, tracers, params))
>    1021   return apply_todos(env_trace_todo(), outs)
>    1022 
> 
> ~/Codes/JAX_source/jax/jax/interpreters/batching.py in process_call(self, call_primitive, f, tracers, params)
>     146     else:
>     147       f, dims_out = batch_subtrace(f, self.master, dims)
> --> 148       vals_out = call_primitive.bind(f, *vals, **params)
>     149       return [BatchTracer(self, v, d) for v, d in zip(vals_out, dims_out())]
>     150 
> 
> ~/Codes/JAX_source/jax/jax/core.py in _call_bind(processor, post_processor, primitive, f, *args, **params)
>    1018     tracers = map(top_trace.full_raise, args)
>    1019     process = getattr(top_trace, processor)
> -> 1020     outs = map(full_lower, process(primitive, f, tracers, params))
>    1021   return apply_todos(env_trace_todo(), outs)
>    1022 
> 
> ~/Codes/JAX_source/jax/jax/interpreters/ad.py in process_call(self, call_primitive, f, tracers, params)
>     342     name = params.get('name', f.__name__)
>     343     params = dict(params, name=wrap_name(name, 'jvp'))
> --> 344     result = call_primitive.bind(f_jvp, *primals, *nonzero_tangents, **params)
>     345     primal_out, tangent_out = tree_unflatten(out_tree_def(), result)
>     346     return [JVPTracer(self, p, t) for p, t in zip(primal_out, tangent_out)]
> 
> ~/Codes/JAX_source/jax/jax/core.py in _call_bind(processor, post_processor, primitive, f, *args, **params)
>    1018     tracers = map(top_trace.full_raise, args)
>    1019     process = getattr(top_trace, processor)
> -> 1020     outs = map(full_lower, process(primitive, f, tracers, params))
>    1021   return apply_todos(env_trace_todo(), outs)
>    1022 
> 
> ~/Codes/JAX_source/jax/jax/interpreters/batching.py in process_call(self, call_primitive, f, tracers, params)
>     146     else:
>     147       f, dims_out = batch_subtrace(f, self.master, dims)
> --> 148       vals_out = call_primitive.bind(f, *vals, **params)
>     149       return [BatchTracer(self, v, d) for v, d in zip(vals_out, dims_out())]
>     150 
> 
> ~/Codes/JAX_source/jax/jax/core.py in _call_bind(processor, post_processor, primitive, f, *args, **params)
>    1014   if top_trace is None:
>    1015     with new_sublevel():
> -> 1016       outs = primitive.impl(f, *args, **params)
>    1017   else:
>    1018     tracers = map(top_trace.full_raise, args)
> 
> ~/Codes/JAX_source/jax/jax/interpreters/xla.py in _xla_call_impl(fun, device, backend, name, *args)
>     466 
>     467 def _xla_call_impl(fun: lu.WrappedFun, *args, device, backend, name):
> --> 468   compiled_fun = _xla_callable(fun, device, backend, name, *map(arg_spec, args))
>     469   try:
>     470     return compiled_fun(*args)
> 
> ~/Codes/JAX_source/jax/jax/linear_util.py in memoized_fun(fun, *args)
>     219       fun.populate_stores(stores)
>     220     else:
> --> 221       ans = call(fun, *args)
>     222       cache[key] = (ans, fun.stores)
>     223     return ans
> 
> ~/Codes/JAX_source/jax/jax/interpreters/xla.py in _xla_callable(fun, device, backend, name, *arg_specs)
>     516   xla_consts = _map(partial(xb.constant, c), consts)
>     517   xla_args = _xla_callable_args(c, abstract_args, tuple_args)
> --> 518   out_nodes = jaxpr_subcomp(
>     519       c, jaxpr, backend, AxisEnv(nreps, (), ()), xla_consts,
>     520       extend_name_stack(wrap_name(name, 'jit')), *xla_args)
> 
> ~/Codes/JAX_source/jax/jax/interpreters/xla.py in jaxpr_subcomp(c, jaxpr, backend, axis_env, consts, name_stack, *args)
>     344       new_params = check_backend_params(eqn.params, backend)
>     345       rule = initial_style_translations[eqn.primitive]
> --> 346       ans = rule(c, axis_env, extend_name_stack(name_stack, eqn.primitive.name),
>     347                  map(aval, eqn.invars), backend, *in_nodes, **new_params)
>     348     elif eqn.primitive in parallel_translations:
> 
> ~/Codes/JAX_source/jax/jax/lax/lax_control_flow.py in _while_loop_translation_rule(c, axis_env, name_stack, avals, backend, cond_jaxpr, body_jaxpr, cond_nconsts, body_nconsts, *args)
>     288                                    _map(partial(xb.constant, body_c), cond_jaxpr.literals),
>     289                                    extend_name_stack(name_stack, 'body_pred'), *(x + z))
> --> 290     new_z = _map(partial(_pred_bcast_select, body_c, body_pred), new_z, z)
>     291     assert _map(body_c.GetShape, new_z) == _map(body_c.GetShape, z) # no broadcast
>     292   new_carry = xops.Tuple(body_c, list(itertools.chain(x, y, new_z)))
> 
> ~/Codes/JAX_source/jax/jax/util.py in safe_map(f, *args)
>      32   for arg in args[1:]:
>      33     assert len(arg) == n, 'length mismatch: {}'.format(list(map(len, args)))
> ---> 34   return list(map(f, *args))
>      35 
>      36 def unzip2(xys):
> 
> ~/Codes/JAX_source/jax/jax/lax/lax_control_flow.py in _pred_bcast_select(c, pred, x, y)
>     302   y_shape = c.GetShape(y).dimensions()
>     303   assert x_shape == y_shape
> --> 304   assert pred_shape == x_shape[:len(pred_shape)] == y_shape[:len(pred_shape)]
>     305   bcast_pred = xops.BroadcastInDim(pred, x_shape, list(range(len(pred_shape))))
>     306   return xops.Select(bcast_pred, x, y)
> 
> AssertionError: 
```
Good point re JVP vs VJP.

I'm going to mark this down as a bug for now. At the very least it's an uninformative error.
Thank you! Any comment whether it could be related to the issue brought up in #3056  ?That would lead to an understanding of how to fix this. 
#3056 is specific to reverse-mode (i.e. VJPs) so I don't think it's related.
I think this is a bug in our while_loop batching rule (and a while_loop is used in expm, via the fori_loop wrapper)! It's a bit hard to articulate, but I think I see it...

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
@@ -375,7 +375,8 @@ def _while_loop_batching_rule(args, dims, cond_nconsts, cond_jaxpr,
     body_jaxpr_batched, carry_bat_out = batching.batch_jaxpr(
         body_jaxpr, size, batched, instantiate=carry_bat)
     cond_jaxpr_batched, (pred_bat,) = batching.batch_jaxpr(
-        cond_jaxpr, size, cconst_bat + carry_bat, instantiate=False)
+        cond_jaxpr, size, cconst_bat + carry_bat,
+        instantiate=bool(cond_jaxpr.out_avals[0].shape))
     carry_bat_out = _map(partial(operator.or_, pred_bat), carry_bat_out)
     if carry_bat_out == carry_bat:
       break
@@ -389,7 +390,7 @@ def _while_loop_batching_rule(args, dims, cond_nconsts, cond_jaxpr,
   new_consts = [batching.moveaxis(x, d, 0) if d is not batching.not_mapped and d != 0
                 else x for x, d in zip(consts, const_dims)]
   new_init = [batching.broadcast(x, size, 0) if now_bat and not was_bat
-              else batching.moveaxis(x, d, 0) if now_bat else x
+              else batching.moveaxis(x, d, 0) if now_bat and d != 0 else x
               for x, d, was_bat, now_bat in zip(init, init_dims, init_bat, carry_bat)]
 
   outs = while_p.bind(*(new_consts + new_init),


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
@@ -1982,6 +1982,12 @@ def trivial_while(y):
     with api.disable_jit():
       api.vmap(trivial_while)(jnp.array([3.0,4.0]))  # doesn't crash
 
+  def test_vmaps_of_while_loop(self):
+    # https://github.com/google/jax/issues/3164
+    def f(x, n): return lax.fori_loop(0, n, lambda _, x: x + 1, x)
+    x, n = jnp.arange(3), jnp.arange(4)
+    api.vmap(api.vmap(f, (None, 0)), (0, None))(x, n)  # doesn't crash
+
 
 if __name__ == '__main__':
   absltest.main()