Repository: google/jax
Issue: scan+pmap shape bug

@JasperSnoek reported this one:

```python
import jax.numpy as np
from jax import lax, random
from jax import pmap

from jax.lib import xla_bridge
nrep = xla_bridge.device_count()

key = random.PRNGKey(0)

n = 80
a_in = np.linspace(0, 1, n)[:, None]
a = np.dot(a_in, a_in.T) + np.eye(n) * 0.1
b = random.normal(key, (n,)).flatten()

def pmvm(a, b):
  a = a.reshape((nrep, -1, a.shape[1]))
  func = pmap(lambda z: np.dot(z, b))
  return func(a).reshape(b.shape)

for i in range(10):
  b = pmvm(a, b)
# Ok

iters = np.arange(10)
def body(carry, i):
  return pmvm(a, carry), i
res = lax.scan(body, b, iters)
```

It results in some kind of shape error (the exact one depends on how many replicas there are).

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

diff --git a/jax/interpreters/partial_eval.py b/jax/interpreters/partial_eval.py
index 69d4ffbd9..23493de5a 100644
--- a/jax/interpreters/partial_eval.py
+++ b/jax/interpreters/partial_eval.py
@@ -143,6 +143,8 @@ class JaxprTrace(Trace):
     return out_tracers
 
   def post_process_call(self, call_primitive, out_tracers, params):
+    if call_primitive in map_primitives:
+      return self.post_process_map(call_primitive, out_tracers, params)
     jaxpr, consts, env = tracers_to_jaxpr([], out_tracers)
     out_pvs, out_pv_consts = unzip2(t.pval for t in out_tracers)
     out = out_pv_consts + consts
@@ -163,6 +165,31 @@ class JaxprTrace(Trace):
       return out_tracers
     return out, todo
 
+  def post_process_map(self, map_primitive, out_tracers, params):
+    jaxpr, consts, env = tracers_to_jaxpr([], out_tracers)
+    out_pvs_reduced, out_pv_consts = unzip2(t.pval for t in out_tracers)
+    out_pvs = [None if pv is None else _unmapped_aval(params['axis_size'], pv)
+               for pv in out_pvs_reduced]
+    out = out_pv_consts + consts
+    del consts, out_pv_consts
+    master = self.master
+    def todo(x):
+      n = len(jaxpr.outvars)
+      out_pv_consts, consts = x[:n], x[n:]
+      trace = JaxprTrace(master, core.cur_sublevel())
+      const_tracers = map(trace.new_instantiated_const, consts)
+      env_tracers = map(trace.full_raise, env)
+      lifted_jaxpr = closure_convert_jaxpr(jaxpr)
+      bound_subjaxpr = (lifted_jaxpr, (), env_tracers)
+      out_tracers = [JaxprTracer(trace, PartialVal((out_pv, out_pv_const)), None)
+                     for out_pv, out_pv_const in zip(out_pvs, out_pv_consts)]
+      eqn = new_jaxpr_eqn(const_tracers, out_tracers, map_primitive,
+                          (bound_subjaxpr,), params)
+      for t in out_tracers:
+        t.recipe = eqn
+      return out_tracers
+    return out, todo
+
 def _mapped_aval(aval):
   if aval is core.abstract_unit:
     return aval
diff --git a/tests/pmap_test.py b/tests/pmap_test.py
index 51bab0541..83e3ebfcb 100644
--- a/tests/pmap_test.py
+++ b/tests/pmap_test.py
@@ -689,6 +689,27 @@ class PmapTest(jtu.JaxTestCase):
     z = y[0]  # doesn't crash
     self.assertAllClose(z, 2 * x[0], check_dtypes=False)
 
+  def testPostProcessMap(self):
+    # test came from https://github.com/google/jax/issues/1369
+    nrep = xla_bridge.device_count()
+
+    def pmvm(a, b):
+      a = a.reshape((nrep, -1, a.shape[1]))
+      func = pmap(lambda z: np.dot(z, b))
+      return func(a).reshape(b.shape)
+
+    rng = onp.random.RandomState(0)
+    a = rng.randn(80, 80)
+    b = rng.randn(80)
+
+    iters = np.arange(5)
+    def body(carry, i):
+      return pmvm(a, carry), i
+    ans, _ = lax.scan(body, b, iters)
+
+    expected = onp.linalg.matrix_power(a, 5).dot(b)
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
 
 class PmapWithDevicesTest(jtu.JaxTestCase):