# Buggy code
def f_jitted(*args, **kwargs):
    context = (
        getattr(core.thread_local_state.trace_state.trace_stack, "dynamic", None),
        config.x64_enabled,
    )
    # TODO(jblespiau): Move this to C++.
    if (FLAGS.jax_debug_nans or FLAGS.jax_debug_infs) and not _jit_is_disabled():
        device_arrays = cpp_jitted_f(context, *args, **kwargs)
        try:
            xla.check_special(
                xla.xla_call_p,
                [
                    da.device_buffer
                    for da in tree_leaves(device_arrays)
                    if hasattr(da, "device_buffer")
                ],
            )
            return device_arrays
        except FloatingPointError:
            assert (
                FLAGS.jax_debug_nans or FLAGS.jax_debug_infs
            )  # compiled_fun can only raise in this case
            print(
                "Invalid nan value encountered in the output of a C++-jit "
                "function. Calling the de-optimized version."
            )
            return cache_miss(context, *args, **kwargs)[0]  # probably won't return
    elif _jit_is_disabled():
        return cpp_jitted_f(*args, **kwargs)
    else:
        return cpp_jitted_f(context, *args, **kwargs)


# Traceback (RuntimeError)
Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
jax._src.traceback_util.FilteredStackTrace: RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

The stack trace above excludes JAX-internal frames.
The following is the original exception that occurred, unmodified.

--------------------

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
File "tmp.py", line 11, in <module>
f()  # <--- fails
File "tmp.py", line 4, in <lambda>
f = lambda: random.uniform(random.PRNGKey(0), (10,), 'float64', -1, 1)
File "/Users/vanderplas/github/google/jax/jax/_src/random.py", line 365, in uniform
return _uniform(key, shape, dtype, minval, maxval)  # type: ignore
File "/Users/vanderplas/github/google/jax/jax/_src/traceback_util.py", line 139, in reraise_with_filtered_traceback
return fun(*args, **kwargs)
File "/Users/vanderplas/github/google/jax/jax/api.py", line 401, in f_jitted
return cpp_jitted_f(context, *args, **kwargs)
RuntimeError: Invalid argument: Argument does not match host shape or layout of computation parameter 1: want s64[], got s32[]: while running replica 0 and partition 0 of a replicated computation (other replicas may have failed as well).

# Fixed code
def f_jitted(*args, **kwargs):
    # TODO(jblespiau): We can remove `config.x64_enabled` when jaxlib has
    # extension version 4
    context = (
        getattr(core.thread_local_state.trace_state.trace_stack, "dynamic", None),
        config.x64_enabled,
    )
    # TODO(jblespiau): Move this to C++.
    if (FLAGS.jax_debug_nans or FLAGS.jax_debug_infs) and not _jit_is_disabled():
        device_arrays = cpp_jitted_f(context, *args, **kwargs)
        try:
            xla.check_special(
                xla.xla_call_p,
                [
                    da.device_buffer
                    for da in tree_leaves(device_arrays)
                    if hasattr(da, "device_buffer")
                ],
            )
            return device_arrays
        except FloatingPointError:
            assert (
                FLAGS.jax_debug_nans or FLAGS.jax_debug_infs
            )  # compiled_fun can only raise in this case
            print(
                "Invalid nan value encountered in the output of a C++-jit "
                "function. Calling the de-optimized version."
            )
            return cache_miss(context, *args, **kwargs)[0]  # probably won't return
    elif _jit_is_disabled():
        return cpp_jitted_f(*args, **kwargs)
    else:
        return cpp_jitted_f(context, *args, **kwargs)