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. #14087
Job | Run time |
---|---|
3s | |
8s | |
17s | |
11s | |
13s | |
1m 9s | |
1m 18s | |
1m 15s | |
1m 21s | |
1m 10s | |
1m 16s | |
1m 15s | |
9m 36s |