<p dir="auto">I think the issue can be traced down to <code class="notranslate">lax.eq(np.inf, np.nan)</code> returning True, for some reason. I am on a recent master commit (7328689).</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="&gt;&gt;&gt; import jax.numpy as np
&gt;&gt;&gt; from jax.numpy.lax_numpy import isinf
&gt;&gt;&gt; from jax import lax
&gt;&gt;&gt; nan = np.array([np.nan])
&gt;&gt;&gt; print(nan)
 array([nan], dtype=float32)
&gt;&gt;&gt; np.isnan(nan)  # Should be True
 array([False])
&gt;&gt;&gt; np.isinf(nan) # Should be False
 array([True])
&gt;&gt;&gt; lax.eq(np.nan, np.inf)  # Should be False
 array(True)"><pre class="notranslate"><span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-k">import</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span> <span class="pl-k">as</span> <span class="pl-s1">np</span>
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-k">from</span> <span class="pl-s1">jax</span>.<span class="pl-s1">numpy</span>.<span class="pl-s1">lax_numpy</span> <span class="pl-k">import</span> <span class="pl-s1">isinf</span>
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-k">from</span> <span class="pl-s1">jax</span> <span class="pl-k">import</span> <span class="pl-s1">lax</span>
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-s1">nan</span> <span class="pl-c1">=</span> <span class="pl-s1">np</span>.<span class="pl-en">array</span>([<span class="pl-s1">np</span>.<span class="pl-s1">nan</span>])
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-en">print</span>(<span class="pl-s1">nan</span>)
 <span class="pl-en">array</span>([<span class="pl-s1">nan</span>], <span class="pl-s1">dtype</span><span class="pl-c1">=</span><span class="pl-s1">float32</span>)
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-s1">np</span>.<span class="pl-en">isnan</span>(<span class="pl-s1">nan</span>)  <span class="pl-c"># Should be True</span>
 <span class="pl-en">array</span>([<span class="pl-c1">False</span>])
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-s1">np</span>.<span class="pl-en">isinf</span>(<span class="pl-s1">nan</span>) <span class="pl-c"># Should be False</span>
 <span class="pl-en">array</span>([<span class="pl-c1">True</span>])
<span class="pl-c1">&gt;&gt;</span><span class="pl-c1">&gt;</span> <span class="pl-s1">lax</span>.<span class="pl-en">eq</span>(<span class="pl-s1">np</span>.<span class="pl-s1">nan</span>, <span class="pl-s1">np</span>.<span class="pl-s1">inf</span>)  <span class="pl-c"># Should be False</span>
 <span class="pl-en">array</span>(<span class="pl-c1">True</span>)</pre></div>
<p dir="auto">as a result np.nan_to_num, np.nanmean, etc all don't work</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import jax.numpy as np
a = np.zeros(1) / np.zeros(1)
print(a.__array__())
print(np.isnan(a).__array__())"><pre class="notranslate"><code class="notranslate">import jax.numpy as np
a = np.zeros(1) / np.zeros(1)
print(a.__array__())
print(np.isnan(a).__array__())
</code></pre></div>
<p dir="auto"><code class="notranslate">[nan]</code><br>
<code class="notranslate">[False]</code></p>
<p dir="auto">This bug only happens with the CPU-only build of JAX, when I see this warning: "warnings.warn('No GPU found, falling back to CPU.')"</p>