Remove references to deprecated jax aliases - Githubissues

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

Remove references to deprecated jax aliases #683

Closed copybara-service[bot] closed 1 year ago

copybara-service[bot] commented 1 year ago

Remove references to deprecated jax aliases

jax version 0.4.7 deprecated several top-level aliases; these have been raising warnings since March 2023, and will soon be removed. They include

jax.ad -> jax.interpreters.ad
jax.flatten_fun_nokwargs - jax.api_util.flatten_fun_nokwargs
jax.partial_eval -> jax.interpreters.partial_eval
jax.pxla -> jax.interpreters.pxla
jax.xla -> jax.interpreters.xla

This change renames these references in preparation for them to be removed from JAX