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，适用于转换。