Repository: google/jax

Issue:
Debugging nan when using 'pmap'
I was trying to set up error propagation by returning nan values as explained in this [issue](https://github.com/google/jax/issues/4257) and I noticed that to make it work for a _pmap_-ed function I also need to _jit_ it.

So something like this doesn't give me an error:
```python
from jax.config import config
config.update("jax_debug_nans", True)

@jax.pmap
def my_func(x):
  return jnp.nan

x = jnp.zeros((1,))
y = my_func(x)
```

While by adding `jax.jit` I get the error (which is what I want): 
```python
from jax.config import config
config.update("jax_debug_nans", True)

@jax.jit
@jax.pmap
def my_func(x):
  return jnp.nan

x = jnp.zeros((1,))
y = my_func(x)
```

So basically a _pmap_-ed function doesn't behave as a _jit_-ed function or a plain function in regards to debugging nans with `jax_debug_nans`. Adding `jax.jit` is an easy fix but I think it would help if this was documented, maybe [here](https://jax.readthedocs.io/en/latest/notebooks/Common_Gotchas_in_JAX.html#Debugging-NaNs) somewhere. Personally it took me quite a bit of time to figure out what was going on.


Discussion:
We should just make `jax_debug_nans` work with pmap (and also `jax_debug_infs` while we're at it). pmap is generally supposed to work like jit + the parallel mapping, and we also generally don't recommend wrapping pmap in jit since it can result in extra data transfers. I'll work on a fix, thanks for reporting! 

Resolving patch:
diff --git a/jax/interpreters/pxla.py b/jax/interpreters/pxla.py
--- a/jax/interpreters/pxla.py
+++ b/jax/interpreters/pxla.py
@@ -1165,6 +1165,9 @@ def partitioned_sharding_spec(num_partitions: int,
 def execute_replicated(compiled, backend, in_handler, out_handler, *args):
   input_bufs = in_handler(args)
   out_bufs = compiled.execute_sharded_on_local_devices(input_bufs)
+  if xla.needs_check_special():
+    for bufs in out_bufs:
+      xla.check_special("parallel computation", bufs)
   return out_handler(out_bufs)
 
 
diff --git a/jax/interpreters/xla.py b/jax/interpreters/xla.py
--- a/jax/interpreters/xla.py
+++ b/jax/interpreters/xla.py
@@ -357,7 +357,7 @@ def _execute_compiled_primitive(prim, compiled, result_handler, *args):
   device, = compiled.local_devices()
   input_bufs = list(it.chain.from_iterable(device_put(x, device) for x in args if x is not token))
   out_bufs = compiled.execute(input_bufs)
-  check_special(prim, out_bufs)
+  check_special(prim.name, out_bufs)
   return result_handler(*out_bufs)
 
 def _execute_replicated_primitive(prim, compiled, result_handler, *args):
@@ -370,11 +370,13 @@ def _execute_replicated_primitive(prim, compiled, result_handler, *args):
   ]
   return result_handler(*out_bufs)
 
+def needs_check_special():
+  return FLAGS.jax_debug_infs or FLAGS.jax_debug_nans
 
-def check_special(prim, bufs):
-  if FLAGS.jax_debug_infs or FLAGS.jax_debug_nans:
+def check_special(name, bufs):
+  if needs_check_special():
     for buf in bufs:
-      _check_special(prim.name, buf.xla_shape(), buf)
+      _check_special(name, buf.xla_shape(), buf)
 
 def _check_special(name, xla_shape, buf):
   assert not xla_shape.is_tuple()
@@ -845,7 +847,7 @@ def _execute_compiled(compiled: XlaExecutable, avals, handlers, *args):
   device, = compiled.local_devices()
   input_bufs = list(it.chain.from_iterable(device_put(x, device) for x in args if x is not token))
   out_bufs = compiled.execute(input_bufs)
-  check_special(xla_call_p, out_bufs)
+  check_special(xla_call_p.name, out_bufs)
   return [handler(*bs) for handler, bs in zip(handlers, _partition_outputs(avals, out_bufs))]
 
 def _execute_replicated(compiled: XlaExecutable, avals, handlers, *args):
@@ -856,7 +858,7 @@ def _execute_replicated(compiled: XlaExecutable, avals, handlers, *args):
       buf[0] for buf in compiled.execute_sharded_on_local_devices(
           list(zip(*input_bufs)))
   ]
-  check_special(xla_call_p, out_bufs)
+  check_special(xla_call_p.name, out_bufs)
   return [handler(*bs) for handler, bs in zip(handlers, _partition_outputs(avals, out_bufs))]
 
 def _execute_trivial(jaxpr, device: Optional[Device], consts, avals, handlers, *args):


Test patch:
diff --git a/tests/debug_nans_test.py b/tests/debug_nans_test.py
--- a/tests/debug_nans_test.py
+++ b/tests/debug_nans_test.py
@@ -18,9 +18,11 @@
 
 import jax
 import numpy as np
+from unittest import SkipTest
 
 from jax import test_util as jtu
 from jax import numpy as jnp
+from jax.experimental import pjit
 
 from jax.config import config
 config.parse_flags_with_absl()
@@ -49,6 +51,12 @@ def testJitComputationNoNaN(self):
     ans = jax.jit(jnp.tanh)(A)
     ans.block_until_ready()
 
+  def testJitComputationNaN(self):
+    A = jnp.array(0.)
+    with self.assertRaises(FloatingPointError):
+      ans = jax.jit(lambda x: 0. / x)(A)
+      ans.block_until_ready()
+
   def testSingleResultPrimitiveNaN(self):
     A = jnp.array(0.)
     with self.assertRaises(FloatingPointError):
@@ -71,6 +79,67 @@ def f(x):
       with self.assertRaisesRegex(FloatingPointError, msg):
         f(1)
 
+  def testPmap(self):
+    f = jax.pmap(lambda x: 0. / x)
+
+    with self.assertRaisesRegex(
+        FloatingPointError,
+        r"invalid value \(nan\) encountered in parallel computation"):
+      ans = f(jnp.array([0.]))
+      ans.block_until_ready()
+
+    if jax.device_count() >= 2:
+      with self.assertRaisesRegex(
+          FloatingPointError,
+          r"invalid value \(nan\) encountered in parallel computation"):
+        ans = f(jnp.array([1., 0.]))
+        ans.block_until_ready()
+
+  def testPmapNoNaN(self):
+    ans = jax.pmap(lambda x: 0. / x)(jnp.array([1.]))
+    ans.block_until_ready()
+
+  @jtu.ignore_warning(message=".*is an experimental.*")
+  def testXmap(self):
+    if not config.omnistaging_enabled:
+      raise SkipTest("xmap requires omnistaging")
+
+    f = jax.experimental.maps.xmap(
+        lambda x: 0. / x,
+        in_axes=['i'],
+        out_axes=['i'],
+        axis_resources={'i': 'x'})
+
+    with jax.experimental.maps.mesh(np.array(jax.local_devices()[:1]), ('x',)):
+      with self.assertRaisesRegex(
+          FloatingPointError,
+          r"invalid value \(nan\) encountered in parallel computation"):
+        ans = f(jnp.array([0.]))
+        ans.block_until_ready()
+
+    if jax.device_count() >= 2:
+      with jax.experimental.maps.mesh(np.array(jax.local_devices()[:2]), ('x',)):
+        with self.assertRaises(FloatingPointError):
+          ans = f(jnp.array([1., 0.]))
+          ans.block_until_ready()
+
+  @jtu.ignore_warning(message=".*is an experimental.*")
+  @jtu.skip_on_devices("cpu", "gpu")
+  def testPjit(self):
+    if jax.device_count() < 2:
+      raise SkipTest("test requires >=2 devices")
+
+    p = jax.experimental.PartitionSpec('x')
+    f = pjit.pjit(lambda x: 0. / x,
+                  in_axis_resources=p,
+                  out_axis_resources=p)
+
+    with jax.experimental.maps.mesh(np.array(jax.local_devices()[:2]), ('x',)):
+      with self.assertRaises(FloatingPointError):
+        ans = f(jnp.array([0., 1.]))
+        ans.block_until_ready()
+
+  # TODO(skye): add parallel inf tests, ideally by factoring out test logic
 
 class DebugInfsTest(jtu.JaxTestCase):