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. #14114
Job | Run time |
---|---|
3s | |
6s | |
23s | |
14s | |
13s | |
1m 6s | |
1m 10s | |
1m 5s | |
1m 3s | |
55s | |
1m 4s | |
1m 19s | |
8m 41s |