jax.extend.linear_util 模块

jax.extend.linear_util 模块#

StoreException

WrappedFun(f, transforms, stores, params, ...)

表示一个函数 f,该函数要应用 transforms

cache(call, *[, explain])

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

merge_linear_aux(aux1, aux2)

transformation

将一个或多个转换添加到 WrappedFun。

transformation_with_aux

将一个或多个带有辅助输出的转换添加到 WrappedFun。

wrap_init(f[, params])

将函数 f 包装为 WrappedFun,适用于转换。