jax.extend.linear_util.cache

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

用于接收 WrappedFun 作为第一个参数的函数的备忘化装饰器。

参数:

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

• explain (Callable | None | None)

返回:

一个 call 的备忘化版本。