Repository: google/jax

Issue:
Run tests with CPU mode when on CPU+GPU version
I have built the package with 

```
python build/build.py --enable_march_native --enable_mkl_dnn --enable_cuda
pip install -e build
pip install -e .
```

I think this should build a CPU + GPU version.

When I tried to run `pytest -n 2 tests examples -W ignore`, some of the tests failed. It appeared to me that some of the functions are only support in CPU mode.

I also tried `JAX_ENABLE_X64=1 JAX_NUM_GENERATED_CASES=100 pytest -n auto tests` and `JAX_ENABLE_X64=0 JAX_NUM_GENERATED_CASES=100 pytest -n auto tests`, but neither of them worked.


Discussion:

Resolving patch:
diff --git a/jax/lax/lax_parallel.py b/jax/lax/lax_parallel.py
--- a/jax/lax/lax_parallel.py
+++ b/jax/lax/lax_parallel.py
@@ -20,8 +20,10 @@
 from jax.core import Primitive
 from jax.interpreters import ad
 from jax.interpreters import parallel
+from jax.interpreters import xla
 from jax.interpreters import pxla
-from jax.util import partial
+from jax.util import partial, unzip2
+from jax.lib import xla_bridge
 
 
 ### parallel traceables
@@ -32,6 +34,12 @@ def psum(x, axis_name):
 def pmax(x, axis_name):
   return pmax_p.bind(x, axis_name=axis_name)
 
+def pmin(x, axis_name):
+  return pmin_p.bind(x, axis_name=axis_name)
+
+def ppermute(x, axis_name, perm):
+  return ppermute_p.bind(x, axis_name=axis_name, perm=perm)
+
 def pswapaxes(x, axis_name, axis):
   """Analogue to `np.swapaxes` involving a hidden axis.
 
@@ -72,43 +80,75 @@ def _unbound_name_error(primitive_name, *args, **kwargs):
 def PmapPrimitive(name):
   prim = Primitive(name)
   prim.def_impl(partial(_unbound_name_error, name))
-  prim.def_abstract_eval(lambda x, *args, **kwargs: x)
+  prim.def_abstract_eval(lambda x, *args, **params: x)
   return prim
 
 
-def _psum_serial_pmap_rule(vals, axes):
+def _allreduce_serial_pmap_rule(reducer, vals, axes):
   val, = vals
   axis, = axes
-  return lax._reduce_sum(val, [axis]), None
+  return reducer(val, [axis]), None
 
-def _psum_transpose_rule(t, axis_name):
-  return [t]
+def _allreduce_translation_rule(prim, c, val, device_groups):
+  dtype = c.GetShape(val).numpy_dtype()
+  scalar = xla_bridge.Shape.array_shape(dtype, ())
+  computation = xla.primitive_computation(prim, scalar, scalar)
+  return c.AllReduce(val, computation, replica_groups=device_groups)
 
-def _psum_parallel_translation_rule(c, val, device_groups):
-  if len(device_groups) > 1:
-    return c.CrossReplicaSum(val, device_groups)
-  else:
-    return c.CrossReplicaSum(val)
 
 psum_p = PmapPrimitive('psum')
-psum_p.def_impl(partial(_unbound_name_error, 'psum'))
-psum_p.def_abstract_eval(lambda x, *args, **kwargs: x)
-parallel.serial_pmap_primitive_rules[psum_p] = _psum_serial_pmap_rule
-pxla.parallel_translation_rules[psum_p] = _psum_parallel_translation_rule
-ad.deflinear(psum_p, _psum_transpose_rule)
 parallel.defreducer(lax.reduce_sum_p, psum_p)
+parallel.serial_pmap_primitive_rules[psum_p] = \
+    partial(_allreduce_serial_pmap_rule, lax._reduce_sum)
+# TODO(mattjj): replace translation rule when we update jaxlib
+# pxla.parallel_translation_rules[psum_p] = \
+#     partial(_allreduce_translation_rule, lax.add_p)
+pxla.parallel_translation_rules[psum_p] = \
+    lambda c, val, device_groups: c.CrossReplicaSum(val, device_groups)
+ad.deflinear(psum_p, lambda t, axis_name: [t])
 
 
-def _pmax_serial_pmap_rule(vals, axes):
-  val, = vals
-  axis, = axes
-  return lax._reduce_max(val, [axis]), None
-
 pmax_p = PmapPrimitive('pmax')
-pmax_p.def_impl(partial(_unbound_name_error, 'pmax'))
-pmax_p.def_abstract_eval(lambda x, *args, **kwargs: x)
-parallel.serial_pmap_primitive_rules[pmax_p] = _pmax_serial_pmap_rule
 parallel.defreducer(lax.reduce_max_p, pmax_p)
+parallel.serial_pmap_primitive_rules[pmax_p] = \
+    partial(_allreduce_serial_pmap_rule, lax._reduce_max)
+pxla.parallel_translation_rules[pmax_p] = \
+    partial(_allreduce_translation_rule, lax.max_p)
+
+
+pmin_p = PmapPrimitive('pmin')
+parallel.defreducer(lax.reduce_min_p, pmin_p)
+parallel.serial_pmap_primitive_rules[pmin_p] = \
+    partial(_allreduce_serial_pmap_rule, lax._reduce_min)
+pxla.parallel_translation_rules[pmin_p] = \
+    partial(_allreduce_translation_rule, lax.min_p)
+
+
+def _ppermute_translation_rule(c, x, device_groups, perm):
+  group_size = len(device_groups[0])
+  if not all(0 <= i < group_size and 0 <= j < group_size for i, j in perm):
+    msg = ("ppermute permutation elements must take on values between 0 and "
+           "the group size {}, but got {}.")
+    raise ValueError(msg.format(group_size, perm))
+  sources, dests = unzip2(perm)
+  if not (len(sources) == len(set(sources)) and len(dests) == len(set(dests))):
+    msg = "ppermute sources and destinations must be unique, got {}."
+    raise ValueError(msg.format(perm))
+
+  full_perm = []
+  for grp in device_groups:
+    grp = list(sorted(grp))
+    full_perm.extend((grp[src], grp[dst]) for src, dst in perm)
+  return c.CollectivePermute(x, full_perm)
+
+def _ppermute_transpose_rule(t, perm, axis_name):
+  sources, dests = unzip2(perm)
+  inverse_perm = zip(dests, srcs)
+  return ppermute(t, axis_name=axis_name, perm=inverse_perm)
+
+ppermute_p = PmapPrimitive('ppermute')
+# ad.deflinear(ppermute_p, _ppermute_transpose_rule)  # TODO(mattjj): test this
+pxla.parallel_translation_rules[ppermute_p] = _ppermute_translation_rule
 
 
 def _pswapaxes_serial_pmap_rule(vals, axes, axis):
diff --git a/jax/lib/xla_bridge.py b/jax/lib/xla_bridge.py
--- a/jax/lib/xla_bridge.py
+++ b/jax/lib/xla_bridge.py
@@ -65,7 +65,7 @@ def _check_jaxlib_version():
 
 FLAGS = flags.FLAGS
 flags.DEFINE_bool('jax_enable_x64',
-                  strtobool(os.getenv('JAX_ENABLE_X64', "False")),
+                  strtobool(os.getenv('JAX_ENABLE_X64', 'False')),
                   'Enable 64-bit types to be used.')
 flags.DEFINE_string(
     'jax_xla_backend', 'xla',
@@ -74,12 +74,11 @@ def _check_jaxlib_version():
     'jax_backend_target', 'local',
     'Either "local" or "rpc:address" to connect to a remote service target.')
 flags.DEFINE_string(
-    'jax_platform_name', '',
-    'Platform name for XLA. The default is to attempt to use a '
-    'GPU if available, but fall back to CPU otherwise. To set '
-    'the platform manually, pass "cpu" for CPU or "gpu" for '
-    'GPU.')
-
+    'jax_platform_name',
+    os.getenv('JAX_PLATFORM_NAME', ''),
+    'Platform name for XLA. The default is to attempt to use a GPU if '
+    'available, but fall back to CPU otherwise. To set the platform manually, '
+    'pass "cpu" for CPU or "gpu" for GPU.')
 
 
 def get_compile_options(num_replicas=None):


Test patch:
diff --git a/jax/test_util.py b/jax/test_util.py
--- a/jax/test_util.py
+++ b/jax/test_util.py
@@ -20,6 +20,7 @@
 import re
 import itertools as it
 import os
+from unittest import SkipTest
 
 from absl.testing import absltest
 from absl.testing import parameterized
@@ -33,15 +34,15 @@
 from .config import flags
 from .util import partial
 from .tree_util import tree_multimap, tree_all, tree_map, tree_reduce
+from .lib import xla_bridge
 
 # lbr tests placeholder
 
 
 FLAGS = flags.FLAGS
 flags.DEFINE_enum(
-    'jax_test_dut',
-    'cpu',
-    enum_values=['cpu', 'gpu', 'tpu'],
+    'jax_test_dut', '',
+    enum_values=['', 'cpu', 'gpu', 'tpu'],
     help=
     'Describes the device under test in case special consideration is required.'
 )
@@ -173,11 +174,11 @@ def skip_on_devices(*disabled_devices):
   def skip(test_method):
     @functools.wraps(test_method)
     def test_method_wrapper(self, *args, **kwargs):
-      device = FLAGS.jax_test_dut
+      device = FLAGS.jax_test_dut or xla_bridge.get_backend().platform
       if device in disabled_devices:
         test_name = getattr(test_method, '__name__', '[unknown test]')
-        return absltest.unittest.skip(
-            '{} not supported on {}.'.format(test_name, device.upper()))
+        raise SkipTest('{} not supported on {}.'
+                       .format(test_name, device.upper()))
       return test_method(self, *args, **kwargs)
     return test_method_wrapper
   return skip
@@ -191,9 +192,8 @@ def test_method_wrapper(self, *args, **kwargs):
       flag_value = getattr(FLAGS, flag_name)
       if flag_value == skip_value:
         test_name = getattr(test_method, '__name__', '[unknown test]')
-        return absltest.unittest.skip(
-            '{} not supported when FLAGS.{} is {}'.format(
-                test_name, flag_name, flag_value))
+        raise SkipTest('{} not supported when FLAGS.{} is {}'
+                       .format(test_name, flag_name, flag_value))
       return test_method(self, *args, **kwargs)
     return test_method_wrapper
   return skip
diff --git a/tests/pmap_test.py b/tests/pmap_test.py
--- a/tests/pmap_test.py
+++ b/tests/pmap_test.py
@@ -280,6 +280,110 @@ def testShardedDeviceTuple(self):
     w = jit(lambda x: list(x)[0])(y)
     self.assertAllClose(w, x, check_dtypes=False)
 
+  @jtu.skip_on_devices("cpu", "gpu")
+  def testCollectivePermute(self):
+    device_count = xla_bridge.device_count()
+    rotation = [(i, i + 1 % device_count) for i in range(device_count)]
+    f = lambda x: lax.ppermute(x, perm=rotation, axis_name='i')
+    f = pmap(f, 'i')
+
+    x = np.arange(4 * device_count).reshape((device_count, 4))
+    ans = f(x)
+    expected = onp.roll(x, shift=1, axis=0)
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+  @jtu.skip_on_devices("cpu", "gpu")
+  def testRule30(self):
+    # This is a test of collective_permute implementing a simple halo exchange
+    # to run a rule 30 simulation: https://en.wikipedia.org/wiki/Rule_30
+    # Halo exchange should be useful in spatially-sharded convolutions and in
+    # other simulations.
+    device_count = xla_bridge.device_count()
+
+    def send_right(x, axis_name):
+      left_perm = [(i, (i + 1) % device_count) for i in range(device_count)]
+      return lax.ppermute(x, perm=left_perm, axis_name=axis_name)
+
+    def send_left(x, axis_name):
+      left_perm = [((i + 1) % device_count, i) for i in range(device_count)]
+      return lax.ppermute(x, perm=left_perm, axis_name=axis_name)
+
+    def update_board(board):
+      left = board[:-2]
+      right = board[2:]
+      center = board[1:-1]
+      return lax.bitwise_xor(left, lax.bitwise_or(center, right))
+
+    @partial(pmap, axis_name='i')
+    def step(board_slice):
+      left, right = board_slice[:1], board_slice[-1:]
+      right, left = send_left(left, 'i'), send_right(right, 'i')
+      enlarged_board_slice = np.concatenate([left, board_slice, right])
+      return update_board(enlarged_board_slice)
+
+    board = onp.zeros(40, dtype=bool)
+    board[board.shape[0] // 2] = True
+    reshaped_board = board.reshape((device_count, -1))
+
+    boards = []
+    def print_board(board):
+      boards.append(''.join('*' if x else ' ' for x in board.ravel()))
+
+    print_board(reshaped_board)
+    for _ in range(20):
+      reshaped_board = step(reshaped_board)
+      print_board(reshaped_board)
+
+    ans = '\n'.join(boards)
+    expected = '\n'.join((
+        '                    *                   ',
+        '                   ***                  ',
+        '                  **  *                 ',
+        '                 ** ****                ',
+        '                **  *   *               ',
+        '               ** **** ***              ',
+        '              **  *    *  *             ',
+        '             ** ****  ******            ',
+        '            **  *   ***     *           ',
+        '           ** **** **  *   ***          ',
+        '          **  *    * **** **  *         ',
+        '         ** ****  ** *    * ****        ',
+        '        **  *   ***  **  ** *   *       ',
+        '       ** **** **  *** ***  ** ***      ',
+        '      **  *    * ***   *  ***  *  *     ',
+        '     ** ****  ** *  * *****  *******    ',
+        '    **  *   ***  **** *    ***      *   ',
+        '   ** **** **  ***    **  **  *    ***  ',
+        '  **  *    * ***  *  ** *** ****  **  * ',
+        ' ** ****  ** *  ******  *   *   *** ****',
+        ' *  *   ***  ****     **** *** **   *   ',
+    ))
+
+    print(ans)
+    self.assertEqual(ans, expected)
+
+  @jtu.skip_on_devices("cpu", "gpu")
+  def testReduceMax(self):
+    f = pmap(lambda x: x - lax.pmax(x, 'i'), axis_name='i')
+
+    shape = (xla_bridge.device_count(), 4)
+    x = onp.arange(prod(shape), dtype=onp.float32).reshape(shape)
+    expected = x - onp.max(x, 0)
+
+    ans = f(x)
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
+  @jtu.skip_on_devices("cpu", "gpu")
+  def testReduceMin(self):
+    f = pmap(lambda x: x - lax.pmin(x, 'i'), axis_name='i')
+
+    shape = (xla_bridge.device_count(), 4)
+    x = onp.arange(prod(shape), dtype=onp.float32).reshape(shape)
+    expected = x - onp.min(x, 0)
+
+    ans = f(x)
+    self.assertAllClose(ans, expected, check_dtypes=False)
+
 
 if __name__ == '__main__':
   absltest.main()