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. #14085
Job | Run time |
---|---|
6s | |
18s | |
7s | |
14s | |
12s | |
1m 1s | |
1m 12s | |
1m 17s | |
58s | |
1m 4s | |
1m 11s | |
1m 18s | |
8m 58s |