jax.extend.linear_util.wrap_init

jax.extend.linear_util.wrap_init(f, params=None, *, debug_info=None)

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

参数:

• f (Callable)

• debug_info (TracingDebugInfo | None | None)

返回类型:

WrappedFun