<p dir="auto">I'm experiencing some type promotion errors after upgrading to <code class="notranslate">numpy==1.23.0</code>. I'm guessing it's probably not fully supported yet---my current JAX version is <code class="notranslate">0.3.13</code>.</p>
<p dir="auto">The following is a basic repro to trigger the error:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax
import jax.numpy as jnp

jax.grad(lambda x: jnp.sum(jax.nn.softplus(x)))(jnp.array([1.0, 2.0, 3.0]))"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</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">jnp</span>

<span class="pl-s1">jax</span>.<span class="pl-en">grad</span>(<span class="pl-k">lambda</span> <span class="pl-s1">x</span>: <span class="pl-s1">jnp</span>.<span class="pl-en">sum</span>(<span class="pl-s1">jax</span>.<span class="pl-s1">nn</span>.<span class="pl-en">softplus</span>(<span class="pl-s1">x</span>)))(<span class="pl-s1">jnp</span>.<span class="pl-en">array</span>([<span class="pl-c1">1.0</span>, <span class="pl-c1">2.0</span>, <span class="pl-c1">3.0</span>]))</pre></div>
<p dir="auto">Filtered stack trace:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File &quot;jax_bug.py&quot;, line 4, in &lt;module&gt;
    jax.grad(lambda x: jnp.sum(jax.nn.softplus(x)))(jnp.array([1.0, 2.0, 3.0]))
  File &quot;jax_bug.py&quot;, line 4, in &lt;lambda&gt;
    jax.grad(lambda x: jnp.sum(jax.nn.softplus(x)))(jnp.array([1.0, 2.0, 3.0]))
  File &quot;/home/sean/.local/lib/python3.8/site-packages/jax/_src/nn/functions.py&quot;, line 66, in softplus
    return jnp.logaddexp(x, 0)
  File &quot;/home/sean/.local/lib/python3.8/site-packages/jax/_src/numpy/ufuncs.py&quot;, line 361, in _logaddexp_jvp
    x1, x2, t1, t2 = _promote_args_inexact(&quot;logaddexp_jvp&quot;, x1, x2, t1, t2)
  File &quot;/home/sean/.local/lib/python3.8/site-packages/jax/_src/numpy/util.py&quot;, line 327, in _promote_args_inexact
    return _promote_shapes(fun_name, *_promote_dtypes_inexact(*args))
  File &quot;/home/sean/.local/lib/python3.8/site-packages/jax/_src/numpy/util.py&quot;, line 262, in _promote_dtypes_inexact
    to_dtype, weak_type = dtypes._lattice_result_type(*args)
KeyError: dtype([('float0', 'V')])"><pre class="notranslate"><code class="notranslate">The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "jax_bug.py", line 4, in &lt;module&gt;
    jax.grad(lambda x: jnp.sum(jax.nn.softplus(x)))(jnp.array([1.0, 2.0, 3.0]))
  File "jax_bug.py", line 4, in &lt;lambda&gt;
    jax.grad(lambda x: jnp.sum(jax.nn.softplus(x)))(jnp.array([1.0, 2.0, 3.0]))
  File "/home/sean/.local/lib/python3.8/site-packages/jax/_src/nn/functions.py", line 66, in softplus
    return jnp.logaddexp(x, 0)
  File "/home/sean/.local/lib/python3.8/site-packages/jax/_src/numpy/ufuncs.py", line 361, in _logaddexp_jvp
    x1, x2, t1, t2 = _promote_args_inexact("logaddexp_jvp", x1, x2, t1, t2)
  File "/home/sean/.local/lib/python3.8/site-packages/jax/_src/numpy/util.py", line 327, in _promote_args_inexact
    return _promote_shapes(fun_name, *_promote_dtypes_inexact(*args))
  File "/home/sean/.local/lib/python3.8/site-packages/jax/_src/numpy/util.py", line 262, in _promote_dtypes_inexact
    to_dtype, weak_type = dtypes._lattice_result_type(*args)
KeyError: dtype([('float0', 'V')])
</code></pre></div>
<p dir="auto">Note from the stacktrace I thought the culprit would be from <code class="notranslate">logaddexp</code>, so I tried:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax
import jax.numpy as jnp

jax.grad(lambda x: jnp.sum(jnp.logaddexp(x, 0)))(jnp.array([1.0, 2.0, 3.0]))"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</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">jnp</span>

<span class="pl-s1">jax</span>.<span class="pl-en">grad</span>(<span class="pl-k">lambda</span> <span class="pl-s1">x</span>: <span class="pl-s1">jnp</span>.<span class="pl-en">sum</span>(<span class="pl-s1">jnp</span>.<span class="pl-en">logaddexp</span>(<span class="pl-s1">x</span>, <span class="pl-c1">0</span>)))(<span class="pl-s1">jnp</span>.<span class="pl-en">array</span>([<span class="pl-c1">1.0</span>, <span class="pl-c1">2.0</span>, <span class="pl-c1">3.0</span>]))</pre></div>
<p dir="auto">which is the exact implementation of <code class="notranslate">jax.nn.softplus</code> and it runs without issue</p>
<p dir="auto">Also note that downgrading to <code class="notranslate">numpy==1.22.3</code> gets rid of the error completely</p>
<p dir="auto">Something in the numpy 1.23 release broke autodiff; in particular, float0 arrays are being converted to bool:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax
import numpy
print(numpy.__version__)
print(jax.grad(lambda x: x+0., allow_int=True)(1).dtype)"><pre class="notranslate"><span class="pl-k">import</span> <span class="pl-s1">jax</span>
<span class="pl-k">import</span> <span class="pl-s1">numpy</span>
<span class="pl-en">print</span>(<span class="pl-s1">numpy</span>.<span class="pl-s1">__version__</span>)
<span class="pl-en">print</span>(<span class="pl-s1">jax</span>.<span class="pl-en">grad</span>(<span class="pl-k">lambda</span> <span class="pl-s1">x</span>: <span class="pl-s1">x</span><span class="pl-c1">+</span><span class="pl-c1">0.</span>, <span class="pl-s1">allow_int</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)(<span class="pl-c1">1</span>).<span class="pl-s1">dtype</span>)</pre></div>
<p dir="auto">Newest numpy:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="1.23.0
bool"><pre class="notranslate"><code class="notranslate">1.23.0
bool
</code></pre></div>
<p dir="auto">previous version:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="1.22.4
[('float0', 'V')]"><pre class="notranslate"><code class="notranslate">1.22.4
[('float0', 'V')]
</code></pre></div>