Colab TPU crashes when printing jax.numpy.array

[MatiasCovarrubias]

Since yersterday, I am getting the weirdest Colab TPU crashes in code that I use all the time. GPUs and CPUs work fine. This is a minimal implementation:

import jax.tools.colab_tpu
jax.tools.colab_tpu.setup_tpu()

import jax
import numpy

print(jax.numpy.array([2,2]) # do not crash
print(numpy.mean(numpy.array([2,2]) # do not crash
print(jax.numpy.mean(jax.numpy.array([2,2])) # crashes
print(jax.numpy.mean(numpy.array([2,2])) # crashes

The log shows the following message:
external/org_tensorflow/tensorflow/stream_executor/cuda/cuda_driver.cc:265] failed call to cuInit: CUDA_ERROR_NO_DEVICE: no CUDA-capable device is detected

I have tried solutions specific to that error (like setting os.environ['CUDA_VISIBLE_DEVICES'] = "0") but it doesn't work. The TPU is also freezing at compilation time, but I wanted to start with this more reproducible error first. Something very weird is happening. Any help appreciated!

[jakevdp]

This is also reported here: googlecolab/colabtools#3009 It's unclear what is going on.