Replace deprecated jax.linear_util.wrap_init with jax.extend.linear_util.wrap_init. - #252
Open
copybara-service[bot] wants to merge 1 commit into
Open
Replace deprecated jax.linear_util.wrap_init with jax.extend.linear_util.wrap_init.#252copybara-service[bot] wants to merge 1 commit into
copybara-service[bot] wants to merge 1 commit into