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. #14081
Job | Run time |
---|---|
2s | |
6s | |
19s | |
13s | |
15s | |
1m 43s | |
1m 26s | |
1m 10s | |
1m 4s | |
14s | |
1m 7s | |
1m 44s | |
9m 23s |