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 

Changed files:
jax/interpreters/partial_eval.py, tests/lax_control_flow_test.py

diff --git a/jax/interpreters/partial_eval.py b/jax/interpreters/partial_eval.py
index 69d4ffbd9..c9a159b52 100644
--- 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 @@ class JaxprTrace(Trace):
     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
 
 
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
index 342096e6a..20b68317b 100644
--- a/tests/lax_control_flow_test.py
+++ b/tests/lax_control_flow_test.py
@@ -981,6 +981,17 @@ class LaxControlFlowTest(jtu.JaxTestCase):
     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()