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

Open

copybara-service[bot] wants to merge 1 commit intomainfrom test_677843398

+6-3

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
learned_optimization authors
committed