jax.extend.linear_util.cache

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

Memoization decorator for functions taking a WrappedFun as first argument.

Parameters:

• call (Callable) – 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.

• explain (Callable | None)

Returns:

A memoized version of call.