Repository: google/jax

Issue:
vmap(pmap) different results than map(pmap)
Likely I'm confused, but I recently tracked a bug in my program to a different output when I run (what I expect to be) the same computation with vmap vs. map.

This was run on TPUv2.

```
import functools
import jax
import jax.numpy as jnp
print(jax.local_device_count())

# In the code below, pts only exists to allow us to execute with pmap.

def map_version(qs, pts):
  return jax.lax.map(lambda x: func(x, pts), qs)

def vmap_version(qs, pts):
  return jax.vmap(func, in_axes=(0, None))(qs, pts)

def func(q, pts):
"""Returns q and also the view of q that the pmap'd lambda gets"""
  q_from_pmap = jax.pmap(lambda x, y: y, in_axes=(0, None))(pts, q)
  return q, q_from_pmap

device_count = 2
pts = jnp.ones(device_count)
qs = jnp.asarray(((0,0), (3,3), (2,2)))

print(f'qs.shape: {qs.shape} qs:\n{qs}')
print(pts.shape: {pts.shape} {pts}')

print(f'map version\n-----------')
q, q_from_pmap = map_version(qs, pts)
print(f'q from func:{q.shape}\n{q}')
print(f'q_from_pmap:{q_from_pmap.shape}\n{q_from_pmap}')

print(f'vmap version\n-----------')
q, q_from_pmap = vmap_version(qs, pts)
print(f'q from func:{q.shape}\n{q}')
print(f'q_from_pmap:{q_from_pmap.shape}\n{q_from_pmap}')
```

Output:
```
8
qs.shape: (3, 2) qs:
[[0 0]
 [3 3]
 [2 2]]
points.shape: (2,) [1. 1.]
map version
-----------
q from func:(3, 2)
[[0 0]
 [3 3]
 [2 2]]
q_from_pmap:(3, 2, 2)
[[[0 0]
  [0 0]]

 [[3 3]
  [3 3]]

 [[2 2]
  [2 2]]]
vmap version
-----------
q from func:(3, 2)
[[0 0]
 [3 3]
 [2 2]]
q_from_pmap:(2, 2, 3)
[[[0 3 2]
  [0 3 2]]

 [[0 3 2]
  [0 3 2]]]
```

The map version does exactly what I would expect. The vmap version does not.

If I had to guess, it would be that vmap does some sort of virtual slicing, but that slicing is not maintained once qs is passed through to the pmapped lambda. The pmapped lambda seems to get all the rows of qs, when I would expect it to only get one row per batch.

Please let me know if I'm abusing vmap here, or if it's a bug.


**Update**

If we transpose qs in vmap_version, everything works:
```
def vmap_version(qs, pts):
  return jax.vmap(func, in_axes=(0, None))(qs.T, pts)
```
This is true for my actual program as well, which is doing quite a bit more sophisticated work than this example. Doesn't feel right. :)


Discussion:
I think perhaps in #1959 (and the follow-up fix #2828) we neglected to handle `mapped_invars` correctly for `vmap`-of-`pmap`. That is, BatchTracer.process_map still assumes all arguments are mapped over by a map primitive (i.e. by `pmap`).

cc @gnecula @jekbradbury 

Resolving patch:
diff --git a/jax/interpreters/batching.py b/jax/interpreters/batching.py
--- a/jax/interpreters/batching.py
+++ b/jax/interpreters/batching.py
@@ -161,10 +161,12 @@ def process_map(self, map_primitive, f: lu.WrappedFun, tracers, params):
     if all(dim is not_mapped for dim in dims):
       return map_primitive.bind(f, *vals, **params)
     else:
+      mapped_invars = params['mapped_invars']
       size, = {x.shape[d] for x, d in zip(vals, dims) if d is not not_mapped}
-      vals = [moveaxis(x, d, 1) if d is not not_mapped and d != 1 else x
-              for x, d in zip(vals, dims)]
-      dims = tuple(not_mapped if d is not_mapped else 0 for d in dims)
+      vals = [moveaxis(x, d, 1) if d == 0 and mapped_invar else x
+              for x, d, mapped_invar in zip(vals, dims, mapped_invars)]
+      dims = tuple(not_mapped if d is not_mapped else max(0, d - mapped_invar)
+                   for d, mapped_invar in zip(dims, mapped_invars))
       f, dims_out = batch_subtrace(f, self.master, dims)
       vals_out = map_primitive.bind(f, *vals, **params)
       dims_out = tuple(d + 1 if d is not not_mapped else d for d in dims_out())


Test patch:
diff --git a/tests/pmap_test.py b/tests/pmap_test.py
--- a/tests/pmap_test.py
+++ b/tests/pmap_test.py
@@ -854,6 +854,29 @@ def s(keys):
     ans = s(keys)  # doesn't crash
     self.assertEqual(ans.shape, (13, N_DEVICES))
 
+  def testVmapOfPmap3(self):
+    # https://github.com/google/jax/issues/3399
+    device_count = xla_bridge.device_count()
+    if device_count < 2:
+      raise SkipTest("test requires at least two devices")
+
+    def map_version(qs, pts):
+      return jax.lax.map(lambda x: func(x, pts), qs)
+
+    def vmap_version(qs, pts):
+      return jax.vmap(func, in_axes=(0, None))(qs, pts)
+
+    def func(q, pts):
+      q_from_pmap = jax.pmap(lambda x, y: y, in_axes=(0, None))(pts, q)
+      return q, q_from_pmap
+
+    pts = jnp.ones(device_count)
+    qs = jnp.asarray(((0,0), (3,3), (2,2)))
+
+    _, expected = map_version(qs, pts)
+    _, ans = vmap_version(qs, pts)
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
   def testVmapOfPmapNonLeadingAxis(self):
     device_count = xla_bridge.device_count()
     f0 = lambda x: x
@@ -1210,7 +1233,6 @@ def testPsumOnBooleanDtype(self):
       out = pmap(lambda x: jax.lax.pmean(x, 'i'), 'i')(x)
       self.assertEqual(list(out), [1])
 
-
 class PmapWithDevicesTest(jtu.JaxTestCase):
 
   def testAllDevices(self):