I encountered a strange issue in which if you load both tensorflow and jax in the notebook, one of them fails to work. I created this notebook to reproduce this issue.
- I am using the free tier. Maybe premium users don’t get this error.
- The code works for CPU and GPU runtime.
- I did not test it on the deprecated TPU v1 runtime.
To reproduce this issue:
- Load the notebook
- Connect to the TPU v2 runtime.
- Execute the cells from top to bottom in sequence.
The 3rd cell would report an ioctl failed error. However, if you restart the session and skip the 2nd cell, which loads the jax lib, the 3rd cell works. In this state, if you go back and run the 2nd cell, you would get an ioctl failed error in this cell.
It feels like one of the libs is trying to hold the TPU device exclusively.