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

mainfrom

Commits on Sep 1, 2023

Replace deprecated jax.linear_util.wrap_init with jax.extend.linear_util.wrap_init.

suryabhupa
authored and
copybara-github
committed