Replace `jax.linear_util` -> `jax.extend.linear_util` to suppress `DeprecationWarning`s

google-deepmind / dm-haiku

JAX-based neural network library

https://dm-haiku.readthedocs.io

Apache License 2.0

2.91k stars 231 forks source link

Open copybara-service[bot] opened 1 year ago

copybara-service[bot] commented 1 year ago

Replace jax.linear_util -> jax.extend.linear_util to suppress DeprecationWarnings

Also suppressed a pytype error relating to conditional tree import.