<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
from jax.experimental import enable_x64

@jax.jit
def f(x):
    with enable_x64():
        return x.view(jnp.int64)

f(jnp.arange(10).reshape(5,2))"><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">from</span> <span class="pl-s1">jax</span>.<span class="pl-s1">experimental</span> <span class="pl-k">import</span> <span class="pl-s1">enable_x64</span>

<span class="pl-en">@<span class="pl-s1">jax</span>.<span class="pl-s1">jit</span></span>
<span class="pl-k">def</span> <span class="pl-en">f</span>(<span class="pl-s1">x</span>):
    <span class="pl-k">with</span> <span class="pl-en">enable_x64</span>():
        <span class="pl-k">return</span> <span class="pl-s1">x</span>.<span class="pl-en">view</span>(<span class="pl-s1">jnp</span>.<span class="pl-s1">int64</span>)

<span class="pl-en">f</span>(<span class="pl-s1">jnp</span>.<span class="pl-en">arange</span>(<span class="pl-c1">10</span>).<span class="pl-en">reshape</span>(<span class="pl-c1">5</span>,<span class="pl-c1">2</span>))</pre></div>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
main.py in &lt;module&gt;
----&gt; 1 pack_fragments(jnp.array(-10.01), jnp.array(-1))

    [... skipping hidden 14 frame]

.../python3.9/site-packages/jax/_src/util.py in safe_map(f, *args)
     42   for arg in args[1:]:
     43     assert len(arg) == n, 'length mismatch: {}'.format(list(map(len, args)))
---&gt; 44   return list(map(f, *args))
     45 
     46 def unzip2(xys):

RuntimeError: INVALID_ARGUMENT: Binary op shift-left with different element types: u64[1,2] and u32[1,2]."><pre class="notranslate"><code class="notranslate">---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
main.py in &lt;module&gt;
----&gt; 1 pack_fragments(jnp.array(-10.01), jnp.array(-1))

    [... skipping hidden 14 frame]

.../python3.9/site-packages/jax/_src/util.py in safe_map(f, *args)
     42   for arg in args[1:]:
     43     assert len(arg) == n, 'length mismatch: {}'.format(list(map(len, args)))
---&gt; 44   return list(map(f, *args))
     45 
     46 def unzip2(xys):

RuntimeError: INVALID_ARGUMENT: Binary op shift-left with different element types: u64[1,2] and u32[1,2].
</code></pre></div>
<p dir="auto">minimal repro:</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="from jax import jit, experimental

@jit
def f(a):
  with experimental.enable_x64():
    return 1 + a
f(1)"><pre class="notranslate"><span class="pl-k">from</span> <span class="pl-s1">jax</span> <span class="pl-k">import</span> <span class="pl-s1">jit</span>, <span class="pl-s1">experimental</span>

<span class="pl-en">@<span class="pl-s1">jit</span></span>
<span class="pl-k">def</span> <span class="pl-en">f</span>(<span class="pl-s1">a</span>):
  <span class="pl-k">with</span> <span class="pl-s1">experimental</span>.<span class="pl-en">enable_x64</span>():
    <span class="pl-k">return</span> <span class="pl-c1">1</span> <span class="pl-c1">+</span> <span class="pl-s1">a</span>
<span class="pl-en">f</span>(<span class="pl-c1">1</span>)</pre></div>
<div class="highlight highlight-text-python-traceback notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="RuntimeError: Invalid argument: Binary op add with different element types: s64[] and s32[]."><pre class="notranslate"><span class="pl-en">RuntimeError</span>: <span class="pl-s">Invalid argument: Binary op add with different element types: s64[] and s32[].</span></pre></div>