<p dir="auto">The NaN debugging feature works with JIT (good), but not with pmap (bad).  Please make it work with pmap.</p>
<p dir="auto">When I ran the following script:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="from functools import partial
import jax
import jax.numpy as jnp


@partial(jax.pmap, axis_name='i')
def nan_with_pmap(x):
  return jnp.log(-x)


print(nan_with_pmap(jnp.array([42.])))"><pre class="notranslate"><code class="notranslate">from functools import partial
import jax
import jax.numpy as jnp


@partial(jax.pmap, axis_name='i')
def nan_with_pmap(x):
  return jnp.log(-x)


print(nan_with_pmap(jnp.array([42.])))
</code></pre></div>
<p dir="auto">using</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="JAX_DEBUG_NANS=true python script.py"><pre class="notranslate"><code class="notranslate">JAX_DEBUG_NANS=true python script.py
</code></pre></div>
<p dir="auto">I expected it to crash at the point where the nan was produced, but instead the script prints:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="[nan]"><pre class="notranslate"><code class="notranslate">[nan]
</code></pre></div>
<p dir="auto">Would it be possible to add nan detection to <code class="notranslate">pmap</code> as well?</p>