jax.extend.linear_util 模块

StoreException
WrappedFun(f, f_transformed, transforms, ...)	表示要应用 transforms 的函数 f。
cache(call, *[, explain])	用于将 WrappedFun 作为第一个参数的函数的记忆化装饰器。
merge_linear_aux(aux1, aux2)
变换
带辅助信息的变换
wrap_init(f[, 参数, 调试信息])	将函数 f 封装为 WrappedFun，适用于变换。