Repository: google/jax

Issue:
systematic tests for vmap-of-pmap
We've had some bugs pop up in vmap-of-pmap, with users running into them rather than us catching them up front. We have several tests, but they're all "manual" and don't exercise all possible corner cases. We should make really systematic vmap-of-pmap tests so we're confident things work and we never regress things.


Discussion:

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/lax_vmap_test.py b/tests/lax_vmap_test.py
--- a/tests/lax_vmap_test.py
+++ b/tests/lax_vmap_test.py
@@ -23,12 +23,12 @@
 
 import numpy as onp
 
-import jax
 from jax import api
 from jax import dtypes
 from jax import lax
 from jax import test_util as jtu
 from jax.lib import xla_client
+from jax.util import safe_map, safe_zip
 
 from tests.lax_test import (all_dtypes, CombosWithReplacement,
                             compatible_shapes, default_dtypes, float_dtypes,
@@ -38,6 +38,10 @@
 config.parse_flags_with_absl()
 FLAGS = config.FLAGS
 
+map, unsafe_map = safe_map, map
+zip, unsafe_zip = safe_zip, zip
+
+
 def all_bdims(*shapes):
   bdims = (itertools.chain([cast(Optional[int], None)],
                            range(len(shape) + 1)) for shape in shapes)
@@ -56,17 +60,15 @@ def slicer(x, bdim):
     return lambda i: lax.index_in_dim(x, i, bdim, keepdims=False)
 
 def args_slicer(args, bdims):
-  slicers = list(map(slicer, args, bdims))
+  slicers = map(slicer, args, bdims)
   return lambda i: [sl(i) for sl in slicers]
 
 class LaxVmapTest(jtu.JaxTestCase):
 
   def _CheckBatching(self, op, bdim_size, bdims, shapes, dtypes, rng,
                      rtol=None, atol=None):
-    batched_shapes = list(jax.util.safe_map(partial(add_bdim, bdim_size),
-                                            bdims, shapes))
-    args = [rng(shape, dtype)
-            for shape, dtype in jax.util.safe_zip(batched_shapes, dtypes)]
+    batched_shapes = map(partial(add_bdim, bdim_size), bdims, shapes)
+    args = [rng(shape, dtype) for shape, dtype in zip(batched_shapes, dtypes)]
     args_slice = args_slicer(args, bdims)
     ans = api.vmap(op, bdims)(*args)
     expected = onp.stack([op(*args_slice(i)) for i in range(bdim_size)])
@@ -642,7 +644,7 @@ def testBroadcastShapesReturnsPythonInts(self):
       # Note also that we chose 3 * 5 * 3 * 5 such that it fits in the range of
       # values a bfloat16 can represent exactly to avoid ties.
       for dtype, rng_factory in itertools.chain(
-        zip(float_dtypes + int_dtypes, itertools.repeat(jtu.rand_unique_int)))))
+        unsafe_zip(float_dtypes + int_dtypes, itertools.repeat(jtu.rand_unique_int)))))
   def testTopK(self, shape, dtype, k, bdims, rng_factory):
     rng = rng_factory(self.rng())
     # _CheckBatching doesn't work with tuple outputs, so test outputs separately.
diff --git a/tests/pmap_test.py b/tests/pmap_test.py
--- a/tests/pmap_test.py
+++ b/tests/pmap_test.py
@@ -15,6 +15,7 @@
 
 from concurrent.futures import ThreadPoolExecutor
 from functools import partial
+import itertools as it
 import os
 from random import shuffle
 from unittest import SkipTest
@@ -37,6 +38,9 @@
 from jax.interpreters import pxla
 from jax.interpreters import xla
 
+from tests.lax_test import compatible_shapes
+from tests.lax_vmap_test import all_bdims, add_bdim, args_slicer
+
 from jax.config import config
 config.parse_flags_with_absl()
 
@@ -856,10 +860,6 @@ def s(keys):
 
   def testVmapOfPmap3(self):
     # https://github.com/google/jax/issues/3399
-
-    # TODO(mattjj): re-enable
-    raise SkipTest("temporarily skipping test while debugging others")
-
     device_count = xla_bridge.device_count()
     if device_count < 2:
       raise SkipTest("test requires at least two devices")
@@ -1237,6 +1237,36 @@ def testPsumOnBooleanDtype(self):
       out = pmap(lambda x: jax.lax.pmean(x, 'i'), 'i')(x)
       self.assertEqual(list(out), [1])
 
+class VmapOfPmapTest(jtu.JaxTestCase):
+
+  @parameterized.named_parameters(jtu.cases_from_list(
+      {"testcase_name": f"{shapes}_{vmap_bdims}_{pmap_bdims}",
+       "shapes": shapes, "vmap_bdims": vmap_bdims, "pmap_bdims": pmap_bdims}
+      for shape_group in compatible_shapes
+      for num_args in range(1, 4)
+      for shapes in it.combinations_with_replacement(shape_group, num_args)
+      for vmap_bdims in all_bdims(*shapes)
+      for pmap_bdims in it.product([0, None], repeat=num_args)
+      if not all(bd is None for bd in pmap_bdims)
+  ))
+  def testVmapOfPmap(self, shapes, vmap_bdims, pmap_bdims):
+    vmapped_size = 3
+    pmapped_size = xla_bridge.device_count()
+
+    rng = jtu.rand_default(self.rng())
+
+    def fun(*args):
+      return sum(args)
+
+    final_shapes = map(partial(add_bdim, vmapped_size), vmap_bdims,
+                       map(partial(add_bdim, pmapped_size), pmap_bdims, shapes))
+
+    args = [rng(shape, jnp.float32) for shape in final_shapes]
+    args_slice = args_slicer(args, vmap_bdims)
+    ans = vmap(pmap(fun, in_axes=pmap_bdims), vmap_bdims)(*args)
+    expected = np.stack([fun(*args_slice(i)) for i in range(vmapped_size)])
+    self.assertAllClose(ans, expected)
+
 
 class PmapWithDevicesTest(jtu.JaxTestCase):