Repository: google/jax

Issue:
Bug with shape inference in grad + jit + scan
Encountered with new pip install (jax 0.1.45, jaxlib 0.1.28)

```python
import jax
import jax.numpy as np

    
def loop_stuff(init_x, func, num_iter):
    
    def body(args):
        i, x_new, _ = args
        x = init_x
        x_old = x

        x = func(i, x_old)
        i_new = i + 1
        return i_new, x_new, x_old

    init_vals = (0,
                 init_x,
                 np.zeros_like(init_x))

    def scan_step(args, idx):
        del idx
        return body(args), None

    (iterations, sol, prev_sol), _ = jax.lax.scan(
        f=scan_step,
        init=init_vals,
        xs=np.arange(num_iter),
    )
    return iterations, sol, prev_sol


max_steps = 10


def step(i, x):
    return x - 1

init_x = np.zeros(())

def run_scan(x):
    return loop_stuff(init_x=x, func=step, num_iter=max_steps)

@jax.jit
def jit_run_scan(x):
    return loop_stuff(init_x=x, func=step, num_iter=max_steps)

jax.grad(lambda x: run_scan(x)[1])(init_x)  # runs fine
jax.grad(lambda x: jit_run_scan(x)[1])(init_x)  # fails mysteriously
```

The error:
`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[], (), (), (), (), (), f32[], (), ())) -> pred[]; body: (parameter: (s32[], (), (), (), (), (), f32[], (), ())) -> (s32[], (), (), (), (), (), f32[], f32[], ()); init: (s32[], (), (), (), (), (), f32[], (), ())..`

@mattjj 


Discussion:
Thanks for raising this. Here's a smaller repro:

```python
import jax
import jax.numpy as np

@jax.jit
def jit_run_scan(x):
  def fun(carry, _):
    x, _ = carry
    return (2 * x, 0.), None
  (x, _), _ = jax.lax.scan(fun, (x, 0.), np.arange(3))
  return x

jax.grad(lambda x: jit_run_scan(x))(0.)
```

I think the issue is a literal `0.` in the carry output. I have a feeling we fixed this before, or some other version of it...
It's a strange beast. I'm not sure it is an issue of literal `0.`, but I could be missing something subtle. If you drop third carry term in my example, everything seems to work fine, e.g.,
```python
def loop_stuff(init_x, func, num_iter):
    
    def body(args):
        i, x_new = args
        x = init_x
        x_old = x

        x = func(i, x_old)
        i_new = i + 1
        return i_new, x_new

    init_vals = (0,
                 init_x)

    def scan_step(args, idx):
        del idx
        return body(args), None

    (iterations, sol), _ = jax.lax.scan(
        f=scan_step,
        init=init_vals,
        xs=np.arange(num_iter),
    )
    return iterations, sol
```

You can also rewrite my example without any zeros and still cause the error:
```python
import jax
import jax.numpy as np

    
def loop_stuff(init_x, func, num_iter):
    
    def body(args):
        i, x_new, _ = args
        x = init_x
        x_old = x

        x = func(i, x_old)
        i_new = i + 1
        return i_new, x_new, x_old

    init_vals = (1,
                 init_x,
                 init_x)

    def scan_step(args, idx):
        del idx
        return body(args), None

    (iterations, sol, prev_sol), _ = jax.lax.scan(
        f=scan_step,
        init=init_vals,
        xs=np.arange(num_iter),
    )
    return iterations, sol, prev_sol


max_steps = 10


def step(i, x):
    return x - 1

init_x = np.ones(()) + 1

def run_scan(x):
    return loop_stuff(init_x=x, func=step, num_iter=max_steps)

@jax.jit
def jit_run_scan(x):
    return loop_stuff(init_x=x, func=step, num_iter=max_steps)

jax.grad(lambda x: run_scan(x)[1])(init_x)  # runs fine
jax.grad(lambda x: jit_run_scan(x)[1])(init_x)  # fails mysteriously
```

From playing around with it and the error message, my initial guess was that this occurs when you have two copies of the same tensor as initial carry. The error message makes me think that when the type of one is determined, it affect the other at a time when some part of the code doesn't expect it to change. I haven't sifted through how the backend is implemented, but would something like that be possible in some form or another?
Also, not sure if it helps but I've encountered another bug with scan in a similar setting. Replacing the loop with a unrolled loop works fine. I've been having trouble boiling it down to something self contained that I can post but here is the error message in case it is related in any way:
```
File "/usr/local/lib/python3.6/dist-packages/jax/api.py", line 397, in value_and_grad_f
g = vjp_py(onp.ones((), dtype=dtype))
File "/usr/local/lib/python3.6/dist-packages/jax/api_util.py", line 62, in apply_jaxtree_fun
ans = fun(*args)
File "/usr/local/lib/python3.6/dist-packages/jax/api.py", line 1012, in out_vjp_packed
return out_vjp(cotangent_in)
File "/usr/local/lib/python3.6/dist-packages/jax/interpreters/ad.py", line 112, in vjp_
_, arg_cts = backward_pass(jaxpr, consts, (), dummy_args, dummy_primal_and_ct)
File "/usr/local/lib/python3.6/dist-packages/jax/interpreters/ad.py", line 186, in backward_pass
cts_out = get_primitive_transpose(eqn.primitive)(ct_in, *invals, **eqn.params)
File "/usr/local/lib/python3.6/dist-packages/jax/lax/lax_control_flow.py", line 723, in _scan_transpose
jaxpr_trans = _move_stuff_and_add_add(jaxpr_lifted_trans)
File "/usr/local/lib/python3.6/dist-packages/jax/lax/lax_control_flow.py", line 761, in _move_stuff_and_add_add
assert CTc_aval == CTc_aval2
AssertionError:
In call to configurable 'run' (<function run at 0x7ff0991b20d0>)
```

I'm happy to share the code that causes this but I am not quite ready to broadcast it publicly. It isn't confidential or proprietary in any way so I'm happy to share it privately if that is of any interest to you and the JAX team.
Yeah, all bugs are certainly of interest. Thanks for helping to find them!

If you could open an issue for it, perhaps just with the error message, and then separately email me the repro, that would work. It's useful to open the issue just to help track/organize things.

Re: this particular issue, I'll check out your other repros too to see if it looks like the same bug or a separate one.
In your new repro, there is a literal `0.` in the loop carry output again, but it's complicated to explain why! Here's how to see it for yourself, after running the new repro script and starting a post-mortem debugger:

```
In [2]: %debug
> /usr/local/google/home/mattjj/miniconda3/lib/python3.7/site-packages/jaxlib/xla_client.py(713)GetShape()
    711
    712   def GetShape(self, operand):
--> 713     return self._builder.GetShape(operand)
    714
    715   def SetOpMetadata(self, op_metadata):

ipdb> up
> /usr/local/google/home/mattjj/packages/jax/jax/interpreters/xla.py(294)jaxpr_subcomp()
    292       raise NotImplementedError(msg.format(eqn.primitive.name))
    293
--> 294     c.GetShape(ans)  # force xla to do shape error checking
    295     out_nodes = xla_destructure(c, ans) if eqn.primitive.multiple_results else [ans]
    296     _map(write, eqn.outvars, out_nodes)

ipdb> up
> /usr/local/google/home/mattjj/packages/jax/jax/interpreters/xla.py(241)jaxpr_computation()
    239   freevars = _map(c.ParameterWithShape, freevar_shapes)
    240   args = _map(c.ParameterWithShape, arg_shapes)
--> 241   out_nodes = jaxpr_subcomp(c, jaxpr, backend, axis_env, consts, freevars, *args)
    242   return c.Build(c.Tuple(*out_nodes))
    243

ipdb> p jaxpr
{ lambda  ;  ; a b c d e f g.
  let h i j k l m n o p = while[ cond_nconsts=0
                                 cond_jaxpr={ lambda  ;  ; a b c d e f g h i.
                                              let j = lt a 10
                                              in [j] }
                                 body_nconsts=0
                                 body_jaxpr={ lambda  ;  ; a b c d e f g h i.
                                              let j = add a 1
                                              in [j, *, *, *, *, *, g, 0.0, *] } ] 0 a b c d e f g *
  in [i, j, k, l, m, n, o, p] }
```

In the line with `[j, *, *, *, *, *, g, 0.0, *]` there's a literal `0.0` which I believe is problematic.

If this issue is like the one I'm recalling, it's just an implementation detail around the `core.unit` placeholders we use in partial evaluation.
Thanks for the explanation! I had forgotten that jax transforms the scan into a while op.

Re: other bug, for good measure I tried to recreate it once more and it seems like I can no longer reproduce it with the newest wheels. I was probably missing a fix when I last tried. I hope I didn't waste any of your time on this one!

Resolving patch:
diff --git a/jax/interpreters/partial_eval.py b/jax/interpreters/partial_eval.py
--- a/jax/interpreters/partial_eval.py
+++ b/jax/interpreters/partial_eval.py
@@ -36,19 +36,17 @@ def identity(x): return x
 
 # A partial value (pval) is modeled as a pair (pv, const), as per
 #   type PVal = (PV, Const)
-#   data PV = NonePV | AbstractPV AbstractValue
+#   data PV = Known | Unknown AbstractValue
 #   type Const = MaybeTraced JaxType
-# where the NonePV arm indicates a known (constant) value, the AbstractPV arm
-# indicates an unknown value.
-# Additionally, when the pv is an AbstractValue, then the const must be unit.
+# where the Known arm, represented by a None, indicates a known (constant) value
+# and the Unknown arm, represented by an AbstractValue instance, indicates an
+# unknown value.
+# When the pv is an AbstractValue, then the const must be unit.
 
 
 class JaxprTrace(Trace):
   def pure(self, val):
-    if type(val) in core.literalable_types and onp.shape(val) == ():
-      return JaxprTracer(self, PartialVal((None, val)), Literal(val))
-    else:
-      return self.new_const(val)
+    return self.new_const(val)
 
   def lift(self, val):
     return self.new_const(val)
@@ -76,8 +74,8 @@ def instantiate_const(self, tracer):
     if isinstance(pv, AbstractValue):
       return tracer
     elif pv is None:
-      if type(tracer.recipe) is Literal:
-        return self.new_instantiated_literal(tracer.recipe.val)
+      if type(const) in core.literalable_types and onp.shape(const) == ():
+        return self.new_instantiated_literal(const)
       else:
         return self.new_instantiated_const(const)
     else:
@@ -392,8 +390,7 @@ def tracers_to_jaxpr(in_tracers, out_tracers):
   env_vars, env_vals = unzip2(env.items())
   const_vars, const_vals = unzip2(consts.items())
   jaxpr = Jaxpr(const_vars, env_vars, invars, list(map(var, out_tracers)), eqns)
-  # core.skip_checks or core.check_jaxpr(jaxpr)
-  core.check_jaxpr(jaxpr)
+  core.skip_checks or core.check_jaxpr(jaxpr)
   return jaxpr, const_vals, env_vals
 
 


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
@@ -981,6 +981,17 @@ def f(carry, _):
     key = random.PRNGKey(0)
     api.grad(lambda c: lax.scan(f, (c, key), onp.ones(3))[0][0])(0.)  # doesn't crash
 
+  def testIssue1361(self):
+    @api.jit
+    def jit_run_scan(x):
+      def fun(carry, _):
+        x, _ = carry
+        return (2 * x, 0.), None
+      (x, _), _ = lax.scan(fun, (x, 0.), np.arange(3))
+      return x
+
+    api.grad(lambda x: jit_run_scan(x))(0.)  # doesn't crash
+
 
 if __name__ == '__main__':
   absltest.main()