jax.extend.linear_util.wrap_init

Contents

jax.extend.linear_util.wrap_init#

jax.extend.linear_util.wrap_init(f, params=None)[source]#

Wraps function f as a WrappedFun, suitable for transformation.

Return type:

WrappedFun