Repository: google/jax
Issue: Shape error from loop that returns a constant

This program:
```
from jax import lax

def body(x):
	a, b = x
	return (7, b + 1)

def cond(x):
	a, b = x
	return b < 10

print(lax.while_loop(cond, body, (33, 4))) 
```

produces:
```
$ python loop.py
/usr/local/google/home/phawkins/p/jax/jax/lib/xla_bridge.py:144: UserWarning: No GPU/TPU found, falling back to CPU.
  warnings.warn('No GPU/TPU found, falling back to CPU.')
Traceback (most recent call last):
  File "loop.py", line 11, in <module>
    print(lax.while_loop(cond, body, (33, 4)))
  File "/usr/local/google/home/phawkins/p/jax/jax/lax/lax_control_flow.py", line 108, in while_loop
    aval_out=aval_out, cond_jaxpr=cond_jaxpr, body_jaxpr=body_jaxpr)
  File "/usr/local/google/home/phawkins/p/jax/jax/core.py", line 75, in bind
    return self.impl(*args, **kwargs)
  File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 51, in apply_primitive
    compiled_fun = xla_primitive_callable(prim, *abstract_args, **kwargs)
  File "/usr/local/google/home/phawkins/p/jax/jax/util.py", line 174, in memoized_fun
    ans = cache[key] = fun(*args, **kwargs)
  File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 57, in xla_primitive_callable
    built_c = primitive_computation(prim, *shapes, **kwargs)
  File "/usr/local/google/home/phawkins/p/jax/jax/util.py", line 174, in memoized_fun
    ans = cache[key] = fun(*args, **kwargs)
  File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 74, in primitive_computation
    raise e
  File "/usr/local/google/home/phawkins/p/jax/jax/interpreters/xla.py", line 70, in primitive_computation
    return c.Build()
  File "/usr/local/google/home/phawkins/p/jax/jax/lib/xla_bridge.py", line 283, in Build
    *args, **kwargs)
  File "/usr/local/google/home/phawkins/.pyenv/versions/3.7.3/lib/python3.7/site-packages/jaxlib/xla_client.py", line 875, in Build
    return Computation(self._builder.Build(), backend=backend)
RuntimeError: Invalid argument: The parameter of condition and body, the result of the body, and init must all have the same shape; got Condition: (parameter: ((s32[], s32[]), (s32[]), (s32[]))) -> pred[]; body: (parameter: ((s32[], s32[]), (s32[]), (s32[]))) -> (((), s32[]), (s32[]), (s32[])); init: ((s32[], s32[]), (s32[]), (s32[]))..: 
```

It looks like the loop logic doesn't handle constants or device arrays in the values returned by the body.

(Reduced reproduction of a bug in optimal control example from Vikas.)

Changed files:
jax/interpreters/xla.py, jax/interpreters/partial_eval.py, tests/lax_control_flow_test.py, jax/interpreters/pxla.py, jax/lax/lax_control_flow.py

diff --git a/jax/interpreters/partial_eval.py b/jax/interpreters/partial_eval.py
index 51cb4dc53..c4f25a2e0 100644
--- a/jax/interpreters/partial_eval.py
+++ b/jax/interpreters/partial_eval.py
@@ -169,7 +169,7 @@ def add_axis_to_aval(size, aval):
 
 
 def partial_eval(f, trace, pvs):
-  f = trace_to_subjaxpr(f, trace.master)
+  f = trace_to_subjaxpr(f, trace.master, False)
   return partial_eval_wrapper(f, tuple(pvs))
 
 
@@ -341,10 +341,11 @@ def abstractify(x):
 def trace_unwrapped_to_jaxpr(fun, pvals, **kwargs):
   return trace_to_jaxpr(lu.wrap_init(fun, kwargs), pvals)
 
-def trace_to_jaxpr(fun, pvals):
+def trace_to_jaxpr(fun, pvals, **kwargs):
   """Traces a function, given abstract inputs, to a jaxpr."""
+  instantiate = kwargs.pop('instantiate', False)
   with new_master(JaxprTrace) as master:
-    fun = trace_to_subjaxpr(fun, master)
+    fun = trace_to_subjaxpr(fun, master, instantiate)
     jaxpr, (out_pval, consts, env) = fun.call_wrapped(pvals)
     assert not env
     del master
@@ -352,12 +353,16 @@ def trace_to_jaxpr(fun, pvals):
   return jaxpr, out_pval, consts
 
 @transformation
-def trace_to_subjaxpr(master, pvals):
+def trace_to_subjaxpr(master, instantiate, pvals):
   assert all([isinstance(pv, PartialVal) for pv in pvals]), pvals
   trace = JaxprTrace(master, core.cur_sublevel())
   in_tracers = map(trace.new_arg, pvals)
   out_tracer = yield in_tracers, {}
   out_tracer = trace.full_raise(out_tracer)
+
+  if instantiate:
+    out_tracer = trace.instantiate_const(out_tracer)
+
   jaxpr, consts, env = tracers_to_jaxpr(in_tracers, out_tracer)
   out_pval = out_tracer.pval
   del trace, in_tracers, out_tracer
@@ -472,7 +477,7 @@ def eval_jaxpr_raw(jaxpr, consts, freevar_vals, *args):
 def compiled_call_impl(fun, *args):
   with new_master(JaxprTrace, True) as master:
     pvals = map(abstractify, args)
-    jaxpr, (pval, consts, env) = trace_to_subjaxpr(fun, master).call_wrapped(pvals)
+    jaxpr, (pval, consts, env) = trace_to_subjaxpr(fun, master, False).call_wrapped(pvals)
     jaxpr_ans = eval_jaxpr_raw(jaxpr, consts, env, *args)
     ans = merge_pvals(jaxpr_ans, pval)
     del master, pvals, pval, consts, env, jaxpr_ans, jaxpr
diff --git a/jax/interpreters/pxla.py b/jax/interpreters/pxla.py
index 1eae1f47f..14d0985f9 100644
--- a/jax/interpreters/pxla.py
+++ b/jax/interpreters/pxla.py
@@ -329,7 +329,7 @@ def _shard_aval(axis_size, aval):
 def parallel_callable(fun, axis_name, axis_size, *avals):
   pvals = [PartialVal((aval, core.unit)) for aval in avals]
   with core.new_master(JaxprTrace, True) as master:
-    jaxpr, (pval, consts, env) = trace_to_subjaxpr(fun, master).call_wrapped(pvals)
+    jaxpr, (pval, consts, env) = trace_to_subjaxpr(fun, master, False).call_wrapped(pvals)
     assert not env
     out = compile_replicated(jaxpr, axis_name, axis_size, consts, *avals)
     compiled, nrep, result_shape = out
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
index 625b3ed5b..2cbbd0405 100644
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -516,7 +516,7 @@ def xla_call_impl(fun, *args):
 def xla_callable(fun, *abstract_args):
   pvals = [pe.PartialVal((aval, core.unit)) for aval in abstract_args]
   with core.new_master(pe.JaxprTrace, True) as master:
-    jaxpr, (pval, consts, env) = pe.trace_to_subjaxpr(fun, master).call_wrapped(pvals)
+    jaxpr, (pval, consts, env) = pe.trace_to_subjaxpr(fun, master, False).call_wrapped(pvals)
     assert not env  # no subtraces here (though cond might eventually need them)
     compiled, result_shape = compile_jaxpr(jaxpr, consts, *abstract_args)
     del master, consts, jaxpr, env
diff --git a/jax/lax/lax_control_flow.py b/jax/lax/lax_control_flow.py
index 1407deef5..2681276c6 100644
--- a/jax/lax/lax_control_flow.py
+++ b/jax/lax/lax_control_flow.py
@@ -69,10 +69,26 @@ def while_loop(cond_fun, body_fun, init_val):
   flat_body_fun, out_tree = pytree_fun_to_jaxtupletree_fun(lu.wrap_init(body_fun), (in_tree,))
   flat_cond_fun, _ = pytree_fun_to_jaxtupletree_fun(lu.wrap_init(cond_fun), (in_tree,))
 
-  pval_flat = lax._abstractify(init_val_flat)
-  cond_jaxpr, _, cond_consts = pe.trace_to_jaxpr(flat_cond_fun, (pval_flat,))
-  body_jaxpr, pval_out, body_consts = pe.trace_to_jaxpr(flat_body_fun, (pval_flat,))
-  aval_out, _ = pval_out
+  carry_pval_flat = carry_aval, _ = lax._abstractify(init_val_flat)
+  cond_jaxpr, cond_pval_out, cond_consts = pe.trace_to_jaxpr(flat_cond_fun, (carry_pval_flat,))
+  body_jaxpr, body_pval_out, body_consts = pe.trace_to_jaxpr(flat_body_fun, (carry_pval_flat,), instantiate=True)
+  carry_aval_out, _ = body_pval_out
+  assert isinstance(carry_aval_out, core.AbstractValue)
+  assert carry_aval == core.lattice_join(carry_aval, carry_aval_out)
+
+  cond_pv, cond_const = cond_pval_out
+  if cond_pv is None:
+    # cond_fun evaluates to a constant, so don't need to generate a while_loop
+    if cond_const:
+      raise ValueError("infinite loop with no effects")
+    else:
+      return init_val
+  else:
+    assert isinstance(cond_pv, core.AbstractValue)
+    if (not isinstance(cond_pv, ShapedArray) or cond_pv.shape
+        or cond_pv.dtype != onp.bool_):
+      msg = "while_loop cond_fun must return a scalar boolean, got {}."
+      raise TypeError(msg.format(cond_pv))
 
   # We don't want to promote literal constants as loop arguments; there are
   # sometimes many of them. We pass tracers as loop arguments, but leave
@@ -97,7 +113,6 @@ def while_loop(cond_fun, body_fun, init_val):
   body_split = split_tracers_and_nontracers(body_jaxpr, body_consts)
   body_jaxpr.constvars, body_nontracer_consts, body_tracer_consts = body_split
 
-
   if out_tree() != in_tree:
     raise TypeError("body_fun input and output must have identical structure")
   out_flat = while_p.bind(
@@ -105,7 +120,7 @@ def while_loop(cond_fun, body_fun, init_val):
       core.pack(cond_tracer_consts), core.pack(body_tracer_consts),
       cond_consts=lax._OpaqueParam(cond_nontracer_consts),
       body_consts=lax._OpaqueParam(body_nontracer_consts),
-      aval_out=aval_out, cond_jaxpr=cond_jaxpr, body_jaxpr=body_jaxpr)
+      aval_out=carry_aval_out, cond_jaxpr=cond_jaxpr, body_jaxpr=body_jaxpr)
   return build_tree(out_tree(), out_flat)
 
 
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
index b6785a66b..ed046fc27 100644
--- a/tests/lax_control_flow_test.py
+++ b/tests/lax_control_flow_test.py
@@ -404,6 +404,20 @@ class LaxControlFlowTest(jtu.JaxTestCase):
             (0, 0), lambda x: (x[0], 0),
             (1, 1), lambda x: x)
 
+  def testIssue649(self):
+    from jax import lax
+
+    def body(x):
+      a, b = x
+      return (7, b + 1)
+
+    def cond(x):
+      a, b = x
+      return b < 10
+
+    out = lax.while_loop(cond, body, (33, 4))
+    self.assertEqual(out, (7, 10))
+
 
 if __name__ == '__main__':
   absltest.main()