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