<p dir="auto">When I try a simple scan loop using Jax in 64bit on a Windows 10 machine, I get the following error.</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="import jax
import jax.numpy as jnp
from jax.config import config
config.update('jax_enable_x64', True)

T = 1.0
init_x = 0.0
init_t = 0.0
dt = T / num_steps
def scan_fun(x, t):
    x += dt
    t += dt
    return x, t
xs = jnp.linspace(init_t, T, num_steps, endpoint=True)
jax.lax.scan(scan_fun, init_x, xs[:, jnp.newaxis])[0]
"><pre class="notranslate"><code class="notranslate">import jax
import jax.numpy as jnp
from jax.config import config
config.update('jax_enable_x64', True)

T = 1.0
init_x = 0.0
init_t = 0.0
dt = T / num_steps
def scan_fun(x, t):
    x += dt
    t += dt
    return x, t
xs = jnp.linspace(init_t, T, num_steps, endpoint=True)
jax.lax.scan(scan_fun, init_x, xs[:, jnp.newaxis])[0]

</code></pre></div>
<p dir="auto">Crashes with the following error</p>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="Traceback (most recent call last):
  File &quot;C:\Program Files\JetBrains\PyCharm 2020.3.3\plugins\python\helpers\pydev\_pydevd_bundle\pydevd_exec2.py&quot;, line 3, in Exec
    exec(exp, global_vars, local_vars)
  File &quot;&lt;string&gt;&quot;, line 13, in &lt;module&gt;
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\traceback_util.py&quot;, line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\control_flow.py&quot;, line 1292, in scan
    out = scan_p.bind(*itertools.chain(consts, in_flat),
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\control_flow.py&quot;, line 1880, in scan_bind
    return core.Primitive.bind(scan_p, *args, **params)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\core.py&quot;, line 284, in bind
    out = top_trace.process_primitive(self, tracers, params)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\core.py&quot;, line 622, in process_primitive
    return primitive.impl(*tracers, **params)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\control_flow.py&quot;, line 1379, in _scan_impl
    return _scan_impl_loop(
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\control_flow.py&quot;, line 1340, in _scan_impl_loop
    _, *outs = while_loop(cond_fun, body_fun, init_val)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\traceback_util.py&quot;, line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\control_flow.py&quot;, line 284, in while_loop
    init_vals, init_avals, body_jaxpr, in_tree, *rest = _create_jaxpr(init_val)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\control_flow.py&quot;, line 269, in _create_jaxpr
    cond_jaxpr, cond_consts, cond_tree = _initial_style_jaxpr(cond_fun, in_tree, init_avals)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\util.py&quot;, line 198, in wrapper
    return cached(bool(config.x64_enabled), *args, **kwargs)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\util.py&quot;, line 191, in cached
    return f(*args, **kwargs)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\control_flow.py&quot;, line 73, in _initial_style_jaxpr
    jaxpr, consts, out_tree = _initial_style_open_jaxpr(fun, in_tree, in_avals)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\util.py&quot;, line 198, in wrapper
    return cached(bool(config.x64_enabled), *args, **kwargs)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\util.py&quot;, line 191, in cached
    return f(*args, **kwargs)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\control_flow.py&quot;, line 68, in _initial_style_open_jaxpr
    jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\interpreters\partial_eval.py&quot;, line 1190, in trace_to_jaxpr_dynamic
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\interpreters\partial_eval.py&quot;, line 1200, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\linear_util.py&quot;, line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\control_flow.py&quot;, line 1324, in cond_fun
    return i &lt; length
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\core.py&quot;, line 516, in __lt__
    def __lt__(self, other): return self.aval._lt(self, other)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\numpy\lax_numpy.py&quot;, line 5256, in deferring_binary_op
    return binary_op(self, other)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\numpy\lax_numpy.py&quot;, line 466, in fn
    return lax_fn(x1, x2)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\lax.py&quot;, line 407, in lt
    return lt_p.bind(x, y)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\core.py&quot;, line 284, in bind
    out = top_trace.process_primitive(self, tracers, params)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\interpreters\partial_eval.py&quot;, line 1062, in process_primitive
    out_avals = primitive.abstract_eval(*avals, **params)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\lax.py&quot;, line 1994, in standard_abstract_eval
    return ShapedArray(shape_rule(*avals, **kwargs), dtype_rule(*avals, **kwargs),
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\lax.py&quot;, line 2070, in naryop_dtype_rule
    _check_same_dtypes(name, False, *aval_dtypes)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\lax.py&quot;, line 6158, in _check_same_dtypes
    raise TypeError(msg.format(name, &quot;, &quot;.join(map(str, types))))
TypeError: lt requires arguments to have the same dtypes, got int64, int32."><pre class="notranslate"><code class="notranslate">Traceback (most recent call last):
  File "C:\Program Files\JetBrains\PyCharm 2020.3.3\plugins\python\helpers\pydev\_pydevd_bundle\pydevd_exec2.py", line 3, in Exec
    exec(exp, global_vars, local_vars)
  File "&lt;string&gt;", line 13, in &lt;module&gt;
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\traceback_util.py", line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\control_flow.py", line 1292, in scan
    out = scan_p.bind(*itertools.chain(consts, in_flat),
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\control_flow.py", line 1880, in scan_bind
    return core.Primitive.bind(scan_p, *args, **params)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\core.py", line 284, in bind
    out = top_trace.process_primitive(self, tracers, params)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\core.py", line 622, in process_primitive
    return primitive.impl(*tracers, **params)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\control_flow.py", line 1379, in _scan_impl
    return _scan_impl_loop(
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\control_flow.py", line 1340, in _scan_impl_loop
    _, *outs = while_loop(cond_fun, body_fun, init_val)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\traceback_util.py", line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\control_flow.py", line 284, in while_loop
    init_vals, init_avals, body_jaxpr, in_tree, *rest = _create_jaxpr(init_val)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\control_flow.py", line 269, in _create_jaxpr
    cond_jaxpr, cond_consts, cond_tree = _initial_style_jaxpr(cond_fun, in_tree, init_avals)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\util.py", line 198, in wrapper
    return cached(bool(config.x64_enabled), *args, **kwargs)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\util.py", line 191, in cached
    return f(*args, **kwargs)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\control_flow.py", line 73, in _initial_style_jaxpr
    jaxpr, consts, out_tree = _initial_style_open_jaxpr(fun, in_tree, in_avals)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\util.py", line 198, in wrapper
    return cached(bool(config.x64_enabled), *args, **kwargs)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\util.py", line 191, in cached
    return f(*args, **kwargs)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\control_flow.py", line 68, in _initial_style_open_jaxpr
    jaxpr, _, consts = pe.trace_to_jaxpr_dynamic(wrapped_fun, in_avals)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\interpreters\partial_eval.py", line 1190, in trace_to_jaxpr_dynamic
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\interpreters\partial_eval.py", line 1200, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\control_flow.py", line 1324, in cond_fun
    return i &lt; length
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\core.py", line 516, in __lt__
    def __lt__(self, other): return self.aval._lt(self, other)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\numpy\lax_numpy.py", line 5256, in deferring_binary_op
    return binary_op(self, other)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\numpy\lax_numpy.py", line 466, in fn
    return lax_fn(x1, x2)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\lax.py", line 407, in lt
    return lt_p.bind(x, y)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\core.py", line 284, in bind
    out = top_trace.process_primitive(self, tracers, params)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\interpreters\partial_eval.py", line 1062, in process_primitive
    out_avals = primitive.abstract_eval(*avals, **params)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\lax.py", line 1994, in standard_abstract_eval
    return ShapedArray(shape_rule(*avals, **kwargs), dtype_rule(*avals, **kwargs),
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\lax.py", line 2070, in naryop_dtype_rule
    _check_same_dtypes(name, False, *aval_dtypes)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\lax.py", line 6158, in _check_same_dtypes
    raise TypeError(msg.format(name, ", ".join(map(str, types))))
TypeError: lt requires arguments to have the same dtypes, got int64, int32.
</code></pre></div>
<p dir="auto">When I try a simple cross product calculation using Jax in 64bit on a Windows 10 machine, I get the following error.</p>
<div class="highlight highlight-source-python notranslate position-relative overflow-auto" dir="auto" data-snippet-clipboard-copy-content="import jax.numpy as jnp
from jax.config import config
config.update('jax_enable_x64', True)

x = jnp.array((1.0,1.0,1.0))
y =  jnp.array((0.0,0.0,0.0))
z = jnp.cross(x,y)"><pre class="notranslate"><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">config</span> <span class="pl-k">import</span> <span class="pl-s1">config</span>
<span class="pl-s1">config</span>.<span class="pl-en">update</span>(<span class="pl-s">'jax_enable_x64'</span>, <span class="pl-c1">True</span>)

<span class="pl-s1">x</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">array</span>((<span class="pl-c1">1.0</span>,<span class="pl-c1">1.0</span>,<span class="pl-c1">1.0</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">0.0</span>,<span class="pl-c1">0.0</span>,<span class="pl-c1">0.0</span>))
<span class="pl-s1">z</span> <span class="pl-c1">=</span> <span class="pl-s1">jnp</span>.<span class="pl-en">cross</span>(<span class="pl-s1">x</span>,<span class="pl-s1">y</span>)</pre></div>
<div class="snippet-clipboard-content notranslate position-relative overflow-auto" data-snippet-clipboard-copy-content="Traceback (most recent call last):
  File &quot;C:\Program Files\JetBrains\PyCharm 2020.3.3\plugins\python\helpers\pydev\_pydevd_bundle\pydevd_exec2.py&quot;, line 3, in Exec
    exec(exp, global_vars, local_vars)
  File &quot;&lt;input&gt;&quot;, line 1, in &lt;module&gt;
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\numpy\lax_numpy.py&quot;, line 3938, in cross
    return _cross(a, b, axisa, axisb, axisc)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\traceback_util.py&quot;, line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\api.py&quot;, line 416, in f_jitted
    return cpp_jitted_f(context, *args, **kwargs)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\api.py&quot;, line 297, in cache_miss
    out_flat = xla.xla_call(
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\core.py&quot;, line 1394, in bind
    return call_bind(self, fun, *args, **params)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\core.py&quot;, line 1385, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\core.py&quot;, line 1397, in process
    return trace.process_call(self, fun, tracers, params)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\core.py&quot;, line 625, in process_call
    return primitive.impl(f, *tracers, **params)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\interpreters\xla.py&quot;, line 586, in _xla_call_impl
    compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\linear_util.py&quot;, line 260, in memoized_fun
    ans = call(fun, *args)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\interpreters\xla.py&quot;, line 662, in _xla_callable
    jaxpr, out_avals, consts = pe.trace_to_jaxpr_final(fun, abstract_args)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\interpreters\partial_eval.py&quot;, line 1220, in trace_to_jaxpr_final
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\interpreters\partial_eval.py&quot;, line 1200, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\linear_util.py&quot;, line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\numpy\lax_numpy.py&quot;, line 3922, in _cross
    a0 = a[..., 0]
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\core.py&quot;, line 552, in __getitem__
    def __getitem__(self, idx): return self.aval._getitem(self, idx)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\numpy\lax_numpy.py&quot;, line 4384, in _rewriting_take
    return _gather(arr, treedef, static_idx, dynamic_idx)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\numpy\lax_numpy.py&quot;, line 4391, in _gather
    indexer = _index_to_gather(shape(arr), idx)  # shared with _scatter_update
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\numpy\lax_numpy.py&quot;, line 4571, in _index_to_gather
    i = _normalize_index(i, x_shape[x_axis]) if normalize_indices else i
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\numpy\lax_numpy.py&quot;, line 4224, in _normalize_index
    lax.lt(index, _constant_like(index, 0)),
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\lax.py&quot;, line 407, in lt
    return lt_p.bind(x, y)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\core.py&quot;, line 284, in bind
    out = top_trace.process_primitive(self, tracers, params)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\interpreters\partial_eval.py&quot;, line 1062, in process_primitive
    out_avals = primitive.abstract_eval(*avals, **params)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\lax.py&quot;, line 1994, in standard_abstract_eval
    return ShapedArray(shape_rule(*avals, **kwargs), dtype_rule(*avals, **kwargs),
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\lax.py&quot;, line 2070, in naryop_dtype_rule
    _check_same_dtypes(name, False, *aval_dtypes)
  File &quot;C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\lax.py&quot;, line 6158, in _check_same_dtypes
    raise TypeError(msg.format(name, &quot;, &quot;.join(map(str, types))))
TypeError: lt requires arguments to have the same dtypes, got int64, int32."><pre class="notranslate"><code class="notranslate">Traceback (most recent call last):
  File "C:\Program Files\JetBrains\PyCharm 2020.3.3\plugins\python\helpers\pydev\_pydevd_bundle\pydevd_exec2.py", line 3, in Exec
    exec(exp, global_vars, local_vars)
  File "&lt;input&gt;", line 1, in &lt;module&gt;
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\numpy\lax_numpy.py", line 3938, in cross
    return _cross(a, b, axisa, axisb, axisc)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\traceback_util.py", line 139, in reraise_with_filtered_traceback
    return fun(*args, **kwargs)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\api.py", line 416, in f_jitted
    return cpp_jitted_f(context, *args, **kwargs)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\api.py", line 297, in cache_miss
    out_flat = xla.xla_call(
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\core.py", line 1394, in bind
    return call_bind(self, fun, *args, **params)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\core.py", line 1385, in call_bind
    outs = primitive.process(top_trace, fun, tracers, params)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\core.py", line 1397, in process
    return trace.process_call(self, fun, tracers, params)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\core.py", line 625, in process_call
    return primitive.impl(f, *tracers, **params)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\interpreters\xla.py", line 586, in _xla_call_impl
    compiled_fun = _xla_callable(fun, device, backend, name, donated_invars,
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\linear_util.py", line 260, in memoized_fun
    ans = call(fun, *args)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\interpreters\xla.py", line 662, in _xla_callable
    jaxpr, out_avals, consts = pe.trace_to_jaxpr_final(fun, abstract_args)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\interpreters\partial_eval.py", line 1220, in trace_to_jaxpr_final
    jaxpr, out_avals, consts = trace_to_subjaxpr_dynamic(fun, main, in_avals)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\interpreters\partial_eval.py", line 1200, in trace_to_subjaxpr_dynamic
    ans = fun.call_wrapped(*in_tracers)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\linear_util.py", line 166, in call_wrapped
    ans = self.f(*args, **dict(self.params, **kwargs))
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\numpy\lax_numpy.py", line 3922, in _cross
    a0 = a[..., 0]
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\core.py", line 552, in __getitem__
    def __getitem__(self, idx): return self.aval._getitem(self, idx)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\numpy\lax_numpy.py", line 4384, in _rewriting_take
    return _gather(arr, treedef, static_idx, dynamic_idx)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\numpy\lax_numpy.py", line 4391, in _gather
    indexer = _index_to_gather(shape(arr), idx)  # shared with _scatter_update
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\numpy\lax_numpy.py", line 4571, in _index_to_gather
    i = _normalize_index(i, x_shape[x_axis]) if normalize_indices else i
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\numpy\lax_numpy.py", line 4224, in _normalize_index
    lax.lt(index, _constant_like(index, 0)),
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\lax.py", line 407, in lt
    return lt_p.bind(x, y)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\core.py", line 284, in bind
    out = top_trace.process_primitive(self, tracers, params)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\interpreters\partial_eval.py", line 1062, in process_primitive
    out_avals = primitive.abstract_eval(*avals, **params)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\lax.py", line 1994, in standard_abstract_eval
    return ShapedArray(shape_rule(*avals, **kwargs), dtype_rule(*avals, **kwargs),
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\lax.py", line 2070, in naryop_dtype_rule
    _check_same_dtypes(name, False, *aval_dtypes)
  File "C:\Users\Adam\anaconda3\envs\tensorflow\lib\site-packages\jax\_src\lax\lax.py", line 6158, in _check_same_dtypes
    raise TypeError(msg.format(name, ", ".join(map(str, types))))
TypeError: lt requires arguments to have the same dtypes, got int64, int32.
</code></pre></div>