DeviceTPU

class composer.devices.DeviceTPU

An extension of Device for TPUs.

When running on TPUVMs, you need to export PJRT_DEVICE=TPU. More details.