jax.extend.linear_util.cache

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

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

参数:

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

• explain (Callable | None)

返回:

call 的记忆化版本。