You signed in with another tab or window. Reload to refresh your session.You signed out in another tab or window. Reload to refresh your session.You switched accounts on another tab or window. Reload to refresh your session.Dismiss alert
This is a follow-up issue of discussion #12245. The solution suggested by a project collaborator crashed the Python interpreter on my TPU VM, but it should work cross-platform.
Seems like callbacks do not work right now on Cloud TPU VM because they are using the older stream_executor runtime but will soon switch to a newer runtime that does support callbacks. I'll monitor and update the issue when callbacks work.
Ah yeah, this should work on Cloud TPU as of jax 0.4.8. I'm gonna close this issue, but please comment or reopen if you find things still aren't working!
Description
This is a follow-up issue of discussion #12245. The solution suggested by a project collaborator crashed the Python interpreter on my TPU VM, but it should work cross-platform.
What jax/jaxlib version are you using?
jax v0.3.17, jaxlab v0.3.15
Which accelerator(s) are you using?
TPU v3-8 with libtpu v1.3.0
Additional System Info
Python 3.8.10, TPU VM on GCP running Ubuntu 20.04 (
Linux t1v-n-e307e167-w-0 5.13.0-1023-gcp #28~20.04.1-Ubuntu SMP Wed Mar 30 03:51:07 UTC 2022 x86_64 x86_64 x86_64 GNU/Linux
)The text was updated successfully, but these errors were encountered: