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. #14094
Job | Run time |
---|---|
6s | |
17s | |
2s | |
10s | |
10s | |
1m 23s | |
1s | |
1m 21s | |
1s | |
1m 7s | |
1s | |
1m 23s | |
6m 2s |