[fix] Replace jax.experimental.host_callback with jax.pure_callback
#1242
| Job | Run time |
|---|---|
| 36m 0s | |
| 35m 28s | |
| 35m 2s | |
| 17m 41s | |
| 17m 53s | |
| 18m 4s | |
| 25m 48s | |
| 27m 5s | |
| 23m 38s | |
| 3h 56m 39s |