Replace deprecated jax.linear_util.wrap_init with jax.extend.linear_util.wrap_init.#722
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.#722copybara-service[bot] wants to merge 1 commit into
copybara-service[bot] wants to merge 1 commit into