jax.extend.linear_util 模块# StoreException WrappedFun(f, f_transformed, transforms, ...) 表示要应用 transforms 的函数 f。 cache(call, *[, explain]) 用于将 WrappedFun 作为第一个参数的函数的备忘化装饰器。 merge_linear_aux(aux1, aux2) 转换 带辅助信息的转换 wrap_init(f[, params, debug_info])