search

google-deepmind / distrax

Apache License 2.0
535 stars 32 forks source link

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

Open copybara-service[bot] opened 1 year ago

copybara-service[bot] commented 1 year ago

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