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.#4225

Merged

copybara-service[bot] merged 1 commit intomainfrom test_678351335

Sep 24, 2024

+26-19

Commits on Sep 24, 2024

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.

dougalm
authored and
Flax Authors
committed