Repository: google/jax

Issue:
invalid jaxpr from composition of soft_pmap, _papply, jit, and jit
The following:
```
@jit
def f(x):
  y = 2 * x
  @jit
  def g(z):
    h, axis_name =_papply(lambda x: x * y)
    return soft_pmap(h, axis_name)(z)
  return g(x)

f(onp.arange(4.).reshape((2, 2)))
```
results in the following invalid jaxpr:
```
{ lambda b ;  ; a.
  let c = reshape[ new_sizes=(1, 2, 2)
                   dimensions=None
                   old_sizes=(2, 2) ] a
      d = mul e 2.0
      f = xla_pmap[ axis_size=1
                    axis_name=<axis 0x7fa985eaa890> ] b c
        { lambda  ; b ; c a.
          let d = gather[ operand_shape=(2, 2)
                          slice_sizes=(1, 2)
                          dimension_numbers=GatherDimensionNumbers(offset_dims=(1, 2), collapsed_slice_dims=(), start_index_map=(0, 1)) ] b c
              e = reshape[ new_sizes=(2, 2)
                           dimensions=None
                           old_sizes=(2, 1, 2) ] d
              f = mul a e
          in f } [  ; d ]
      g = reshape[ new_sizes=(2, 2)
                   dimensions=None
                   old_sizes=(1L, 2, 2) ] f
  in g }
```


Discussion:
Simplified to the following:
```
@jit
def f(x):
  y = 2 * x
  @jit
  def g(z):
    return pmap(lambda x: x * y)(z)
  return g(x)
f(onp.arange(1.).reshape((1, 1)))
```
which results in:
```
{ lambda  ;  ; a.
  let b = mul c 2.0
      d = xla_pmap[ axis_size=1
                    axis_name=<axis 0x7f699c8ea050> ] a
        { lambda  ; c ; a.
          let b = reshape[ new_sizes=(1, 1)
                           dimensions=None
                           old_sizes=(1,) ] a
              d = mul b c
          in d } [  ; b ]
  in d }
```

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
@@ -116,8 +116,7 @@ def process_call(self, call_primitive, f, tracers, params):
     out_pv_const, consts = call_primitive.bind(fun, *in_consts, **params)
     out_pv, jaxpr, env = aux()
     const_tracers = map(self.new_instantiated_const, consts)
-    env_tracers = map(self.full_raise, env)
-    bound_subjaxpr = (jaxpr, const_tracers, env_tracers)
+    bound_subjaxpr = (jaxpr, const_tracers, map(self.full_raise, env))
     eqn = JaxprEqn(tracers, None, call_primitive, (bound_subjaxpr,),
                    False, False, params)
     return JaxprTracer(self, PartialVal((out_pv, out_pv_const)), eqn)
@@ -130,12 +129,11 @@ def process_map(self, map_primitive, f, tracers, params):
     out_pv_reduced, jaxpr, env = aux()
     out_pv = add_axis_to_pv(params['axis_size'], out_pv_reduced)
     const_tracers = map(self.new_instantiated_const, consts)
-    env_tracers = map(self.full_raise, env)
     jaxpr_converted = jaxpr.copy()
     jaxpr_converted.constvars = []
     jaxpr_converted.invars = list(it.chain(jaxpr.constvars, jaxpr.invars))
     invars = tuple(it.chain(const_tracers, tracers))
-    bound_subjaxpr = (jaxpr_converted, (), env)
+    bound_subjaxpr = (jaxpr_converted, (), map(self.full_raise, env))
     eqn = JaxprEqn(invars, None, map_primitive, (bound_subjaxpr,),
                    False, False, params)
     return JaxprTracer(self, PartialVal((out_pv, out_const)), eqn)
@@ -465,6 +463,7 @@ def tracers_to_jaxpr(in_tracers, out_tracer):
     if isinstance(recipe, JaxprEqn):
       eqns.append(eqn_tracer_to_var(var, [var(t)], recipe))
     elif isinstance(recipe, LambdaBinding):
+      assert any(t is in_tracer for in_tracer in in_tracers)
       assert in_tracers, "Lambda binding with no args"
     elif isinstance(recipe, FreeVar):
       env[var(t)] = recipe.val


Test patch: