RuntimeError: Requested backend tpu_driver, but it failed to initialize

[TMcMac]

Hey Everyone,

So I am trying to work my way through a Jax tutorial and I am getting stuck when trying to replicate the very basics of jax.pmap(). The issue I seems to be hitting is that I am missing the driver for using TPUs. Here is the code block I am using:

import jax
import jax.numpy as jnp
impurt numpy as np
import requests
import os

if 'TPU_DRIVER_MODE' not in globals():
url = 'http://' + os.environ['COLAB_TPU_ADDR'].split(':')[0] + ':8475/requestversion/tpu_driver_nightly'
resp = requests.post(url)
TPU_DRIVER_MODE = 1

from jax.config import config
config.FLAGS.jax_xla_backend = "tpu_driver"
config.FLAGS.jax_backend_target = "grpc://" + os.environ['COLAB_TPU_ADDR']
print(config.FLAGS.jax_backend_target)

Then I run jax.devices() and I get:

RuntimeError: Requested backend tpu_driver, but it failed to initialize: Not found: Unable to find driver in registry given worker:

I have the TPU selected from the hardware accelerators so it isn't that. Any ideas what I am missing?

[jakevdp]

Have you seen the colab_tpu.setup_tpu() utility? It takes care of most of this logic for you. Here is an example: https://jax.readthedocs.io/en/latest/jax-101/06-parallelism.html#colab-tpu-setup (or open directly in Colab)

I suspect the issue with your code has to do with using the nightly driver – sometimes that causes problems, which is why we generally pin the driver to a specific release that is known to work: https://github.com/google/jax/blob/f004bcb7b8763451b437202549faaea89f771f41/jax/tools/colab_tpu.py#L34-L37

[TMcMac]

Hey Jake,

Thank you! The link you shared got me back on track with:

import jax

import jax.tools.colab_tpu

jax.tools.colab_tpu.setup_tpu()

jax.devices()