Update libraries to use JAX's limited (and ill-advised) trace-state-querying APIs rather than depending on JAX's deeper internals, which are about to change.

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

Update libraries to use JAX's limited (and ill-advised) trace-state-querying APIs rather than depending on JAX's deeper internals, which are about to change. #795

Closed copybara-service[bot] closed 1 month ago

copybara-service[bot] commented 1 month ago

Update libraries to use JAX's limited (and ill-advised) trace-state-querying APIs rather than depending on JAX's deeper internals, which are about to change.