Repository: google/jax

Issue:
Error compiling while_of_pmap
The test is lax_control_flow_tests.py:test_while_loop_of_pmap. It fails on CPU, GPU, and TPU.

For GPU the failure seems to be 
```
F0331 01:03:30.304149    7266 ir_emitter_unnested.cc:1766] Non-OK-status: body->Accept(&ir_emitter_body) status: Unimplemented: Requested AllReduce not implemented on GPU; replica_count: 8; operand_count: 1; IsCrossReplicaAllReduce: 1; NCCL support: 0; first operand array element-type: F64
```

For TPU the failure is perhaps related but shows up elsewhere: 
```
  File "/build/work/c62b4eb2372c3477aa8032626244c507067f/google3/runfiles/google3/third_party/py/jax/interpreters/xla.py", line 169, in apply_primitive
    compiled_fun = xla_primitive_callable(prim, *map(arg_spec, args), **params)
  File "/build/work/c62b4eb2372c3477aa8032626244c507067f/google3/runfiles/google3/third_party/py/jax/interpreters/xla.py", line 194, in xla_primitive_callable
    compiled = built_c.Compile(compile_options=options, backend=backend)
  File "/build/work/c62b4eb2372c3477aa8032626244c507067f/google3/runfiles/google3/third_party/tensorflow/compiler/xla/python/xla_client.py", line 571, in Compile
    return backend.compile(self.computation, compile_options)
  File "/build/work/c62b4eb2372c3477aa8032626244c507067f/google3/runfiles/google3/third_party/tensorflow/compiler/xla/python/xla_client.py", line 151, in compile
    compile_options.tuple_arguments)
RuntimeError: Invalid argument: Invalid (replica_count,computation_count) pair: (8,1)
```


Discussion:

Resolving patch:
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -185,6 +185,11 @@ def xla_primitive_callable(prim, *arg_specs, **params):
     nreps = initial_style_primitive_replicas(params)
   else:
     nreps = 1
+  if nreps > xb.device_count(backend):
+    msg = ("compiling a primitive computation `{}` that requires {} replicas, "
+           "but only {} XLA devices are available on backend {}.")
+    raise ValueError(msg.format(prim, nreps, xb.device_count(backend),
+                                backend.platform))
   built_c = primitive_computation(prim, AxisEnv(nreps), backend, tuple_args,
                                   *avals, **params)
   options = xb.get_compile_options(


Test patch:
diff --git a/tests/lax_control_flow_test.py b/tests/lax_control_flow_test.py
--- a/tests/lax_control_flow_test.py
+++ b/tests/lax_control_flow_test.py
@@ -1816,25 +1816,41 @@ def arange(n):
     expected = onp.arange(10)
     self.assertAllClose(ans, expected, check_dtypes=False)
 
-  @jtu.skip_on_devices("tpu", "gpu", "cpu")# TODO(mattjj): follow up w/ xla
-  # Issue #2554
   def test_while_loop_of_pmap(self):
     # code from jsnoek@
+
     def body(i, x):
       result = api.pmap(lambda z: lax.psum(np.sin(z), 'i'), axis_name='i')(x)
       return result + x
     f_loop = lambda x: lax.fori_loop(0, 3, body, x)
-    ans = f_loop(np.ones(8))
+    ans = f_loop(np.ones(api.device_count()))
     del body, f_loop
 
     def body2(i, x):
       result = np.broadcast_to(np.sin(x).sum(), x.shape)
       return result + x
     g_loop = lambda x: lax.fori_loop(0, 3, body2, x)
-    expected = g_loop(np.ones(8))
+    expected = g_loop(np.ones(api.device_count()))
 
     self.assertAllClose(ans, expected, check_dtypes=False)
 
+  def test_while_loop_of_pmap_error_message(self):
+
+    def body(i, x):
+      result = api.pmap(lambda z: lax.psum(np.sin(z), 'i'), axis_name='i')(x)
+      return result + x
+    f_loop = lambda x: lax.fori_loop(0, 3, body, x)
+
+    too_big = 2 * api.device_count()
+
+    self.assertRaisesRegex(
+        ValueError,
+        re.escape(
+            "compiling a primitive computation `while` that requires {} "
+            "replicas, but only {} XLA devices are available on backend {}."
+            .format(too_big, api.device_count(), jtu.device_under_test())),
+        lambda: f_loop(np.ones(too_big)))
+
 
 if __name__ == '__main__':
   absltest.main()