<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

def f(x):
    return x + 0.0

def _print_type(x):
    print(x)
    print(x.dtype)
    print(x.dtype is jax.dtypes.float0)

def print_type(x):
    _print_type(x)
    jax.jit(_print_type)(x)
    print(&quot;&quot;)

x = jax.grad(f, allow_int=True)(jnp.array(1))
y = jnp.array(False)

print_type(x)
print_type(y)"><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-k">def</span> <span class="pl-en">f</span>(<span class="pl-s1">x</span>):
    <span class="pl-k">return</span> <span class="pl-s1">x</span> <span class="pl-c1">+</span> <span class="pl-c1">0.0</span>

<span class="pl-k">def</span> <span class="pl-en">_print_type</span>(<span class="pl-s1">x</span>):
    <span class="pl-en">print</span>(<span class="pl-s1">x</span>)
    <span class="pl-en">print</span>(<span class="pl-s1">x</span>.<span class="pl-s1">dtype</span>)
    <span class="pl-en">print</span>(<span class="pl-s1">x</span>.<span class="pl-s1">dtype</span> <span class="pl-c1">is</span> <span class="pl-s1">jax</span>.<span class="pl-s1">dtypes</span>.<span class="pl-s1">float0</span>)

<span class="pl-k">def</span> <span class="pl-en">print_type</span>(<span class="pl-s1">x</span>):
    <span class="pl-en">_print_type</span>(<span class="pl-s1">x</span>)
    <span class="pl-s1">jax</span>.<span class="pl-en">jit</span>(<span class="pl-s1">_print_type</span>)(<span class="pl-s1">x</span>)
    <span class="pl-en">print</span>(<span class="pl-s">""</span>)

<span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">jax</span>.<span class="pl-en">grad</span>(<span class="pl-s1">f</span>, <span class="pl-s1">allow_int</span><span class="pl-c1">=</span><span class="pl-c1">True</span>)(<span class="pl-s1">jnp</span>.<span class="pl-en">array</span>(<span class="pl-c1">1</span>))
<span class="pl-s1">y</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>(<span class="pl-c1">False</span>)

<span class="pl-en">print_type</span>(<span class="pl-s1">x</span>)
<span class="pl-en">print_type</span>(<span class="pl-s1">y</span>)</pre></div>
<p dir="auto">With numpy v1.23.0:</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="False
bool
False
Traced&lt;ShapedArray(void[])&gt;with&lt;DynamicJaxprTrace(level=0/1)&gt;
[('float0', 'V')]
False

False
bool
False
Traced&lt;ShapedArray(bool[])&gt;with&lt;DynamicJaxprTrace(level=0/1)&gt;
bool
False"><pre class="notranslate"><code class="notranslate">False
bool
False
Traced&lt;ShapedArray(void[])&gt;with&lt;DynamicJaxprTrace(level=0/1)&gt;
[('float0', 'V')]
False

False
bool
False
Traced&lt;ShapedArray(bool[])&gt;with&lt;DynamicJaxprTrace(level=0/1)&gt;
bool
False
</code></pre></div>
<p dir="auto">Which has the following problems:</p>
<ul dir="auto">
<li>The "outside jit" dtype of <code class="notranslate">x</code> is <code class="notranslate">bool</code> instead of <code class="notranslate">float0</code>.</li>
<li>The "inside jit" dtype of <code class="notranslate">x</code> and <code class="notranslate">y</code> are different, despite having the same "outside jit" dtype.</li>
<li>The "inside jit" comparison of <code class="notranslate">x.dtype is jax.dtypes.float0</code> returns <code class="notranslate">False</code> despite the "inside jit" dtype of <code class="notranslate">x</code> being <code class="notranslate">float0</code>. (Probably something to do with it secretly being a <code class="notranslate">bool</code> instead.)</li>
</ul>
<p dir="auto">This plays merry havoc will all kinds of other things. I spotted this because it was sneaking past the <code class="notranslate">_check_no_float0s</code> checker when calling <code class="notranslate">jnp.allclose</code>. It also means that attempting to do <code class="notranslate">x == y</code> in the above code will result in an exception being thrown.</p>
<p dir="auto">For reference I've found that numpy v1.21.2 works fine.</p>
<p dir="auto">Possibly this is a numpy bug? (I did check to see that there weren't any <code class="notranslate">float0</code>-related numpy bugs reported.)</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>