[Examples] TPU example fail examples/tpu/tpuvm_mnist.yaml

[Michaelvll]

The following error occurs on our master branch:

sky launch --use-spot examples/tpu/tpuvm_mnist.yaml

(tpuvm_mnist, pid=11260)   File "/home/gcpuser/miniconda3/envs/flax/lib/python3.10/site-packages/jax/scipy/linalg.py", line 18, in <module>
(tpuvm_mnist, pid=11260)     from jax._src.scipy.linalg import (
(tpuvm_mnist, pid=11260)   File "/home/gcpuser/miniconda3/envs/flax/lib/python3.10/site-packages/jax/_src/scipy/linalg.py", line 403, in <module>
(tpuvm_mnist, pid=11260)     @_wraps(scipy.linalg.tril)