jax.extend.linear_util.cache

jax.extend.linear_util.cache(call, *, explain=None)

用于接受 WrappedFun 作为第一个参数的函数的记忆化装饰器。

参数:

• call (Callable) – 一个 Python 可调用对象，它接受 WrappedFun 作为其第一个参数。WrappedFun 上的底层转换和参数用作记忆化缓存键的一部分。

• explain (Callable | None | None)

返回值:

call 的记忆化版本。