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