maxtext on Colab TPU

Question

maxtext on Colab TPU

liutianlin0121 opened this issue a year ago · comments

Hello,

I'm wondering if it is possible to use MaxText with TPU on Colab. Currently, when installing MaxText on Colab and calling jax.devices(), it leads to an XLA error:

XlaRuntimeError: UNKNOWN: TPU initialization failed: No ba16c7433 device found.

Below is the Colab link that demonstrates this issue:

https://colab.research.google.com/drive/115pQR-Ia8DsXnBXbOqGLdhJUpshRyhT-?usp=sharing

Tianlin

Rafi Witten · Answer 1 · Fri Aug 25 2023 05:19:03 GMT+0800 (China Standard Time)

Hey -- can you confirm that your issues is just that jax.devices() doesn't work in this environment? That suggests Jax isn't supported in this environment, possibly due to a misconfiguration.