jax.extend.linear_util 模块# StoreException WrappedFun(f, f_transformed, transforms, ...) 表示要应用 transforms 的函数 f。 cache(call, *, explain) 用于接受 WrappedFun 作为第一个参数的函数的记忆化装饰器。 merge_linear_aux(aux1, aux2) transformation transformation_with_aux wrap_init(f[, params, debug_info])
