# Buggy code
def cache(call: Callable):
    """Memoization decorator for functions taking a WrappedFun as first argument.

    Args:
      call: a Python callable that takes a WrappedFun as its first argument. The
        underlying transforms and params on the WrappedFun are used as part of the
        memoization cache key.

    Returns:
       A memoized version of ``call``.
    """
    fun_caches: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary()
    thread_local: threading.local = _CacheLocalContext()

    def memoized_fun(fun: WrappedFun, *args):
        cache = fun_caches.setdefault(fun.f, {})
        if core.debug_state.check_leaks:
            key = (
                _copy_main_traces(fun.transforms),
                fun.params,
                args,
                bool(FLAGS.jax_enable_x64),
            )
        else:
            key = (fun.transforms, fun.params, args, bool(FLAGS.jax_enable_x64))
        result = cache.get(key, None)
        if result is not None:
            ans, stores = result
            fun.populate_stores(stores)
        else:
            ans = call(fun, *args)
            cache[key] = (ans, fun.stores)

        thread_local.most_recent_entry = weakref.ref(ans)
        return ans

    def _most_recent_entry():
        most_recent_entry = thread_local.most_recent_entry
        if most_recent_entry is not None:
            result = most_recent_entry()
            thread_local.most_recent_entry = None
            return result

    memoized_fun.most_recent_entry = _most_recent_entry  # type: ignore
    memoized_fun.cache_clear = fun_caches.clear  # type: ignore

    return memoized_fun


# 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 cache(call: Callable):
    """Memoization decorator for functions taking a WrappedFun as first argument.

    Args:
      call: a Python callable that takes a WrappedFun as its first argument. The
        underlying transforms and params on the WrappedFun are used as part of the
        memoization cache key.

    Returns:
       A memoized version of ``call``.
    """
    fun_caches: weakref.WeakKeyDictionary = weakref.WeakKeyDictionary()
    thread_local: threading.local = _CacheLocalContext()

    def memoized_fun(fun: WrappedFun, *args):
        cache = fun_caches.setdefault(fun.f, {})
        if core.debug_state.check_leaks:
            key = (
                _copy_main_traces(fun.transforms),
                fun.params,
                args,
                config.x64_enabled,
            )
        else:
            key = (fun.transforms, fun.params, args, config.x64_enabled)
        result = cache.get(key, None)
        if result is not None:
            ans, stores = result
            fun.populate_stores(stores)
        else:
            ans = call(fun, *args)
            cache[key] = (ans, fun.stores)

        thread_local.most_recent_entry = weakref.ref(ans)
        return ans

    def _most_recent_entry():
        most_recent_entry = thread_local.most_recent_entry
        if most_recent_entry is not None:
            result = most_recent_entry()
            thread_local.most_recent_entry = None
            return result

    memoized_fun.most_recent_entry = _most_recent_entry  # type: ignore
    memoized_fun.cache_clear = fun_caches.clear  # type: ignore

    return memoized_fun