求助：Colab中使用Jax/Flax配置TPU失败的问题

我已多方搜索并尝试在Colab中搭配Flax配置TPU，其中一份Notebook是我最接近成功的尝试，但出现了如下报错：

/usr/local/lib/python3.9/dist-packages/jax/tools/colab_tpu.py in setup_tpu(tpu_driver_version)
     36 
     37   if not TPU_DRIVER_MODE:
---> 38     colab_tpu_addr = os.environ['COLAB_TPU_ADDR'].split(':')[0]
     39     url = f'http://{colab_tpu_addr}:8475/requestversion/{tpu_driver_version}'
     40     requests.post(url)

/usr/lib/python3.9/os.py in __getitem__(self, key)
    677         except KeyError:
    678             # raise KeyError with the original key value
--> 679             raise KeyError(key) from None
    680         return self.decodevalue(value)
    681 

KeyError: 'COLAB_TPU_ADDR'

几乎所有相关Notebook都会出现该错误，仅另一份Notebook例外，它的报错如下：

RuntimeError                              Traceback (most recent call last)
<ipython-input-3-1360d59fd674> in <module>
      7 
      8 get_ipython().system('pip install flax diffusers transformers ftfy')
----> 9 jax.devices()
     10 

2 frames
/usr/local/lib/python3.9/dist-packages/jax/_src/xla_bridge.py in devices(backend)
    541     List of Device subclasses.
    542   """
--> 543   return get_backend(backend).devices()
    544 
    545 

/usr/local/lib/python3.9/dist-packages/jax/_src/xla_bridge.py in get_backend(platform)
    483 @lru_cache(maxsize=None)  # don't use util.memoize because there is no X64 dependence.
    484 def get_backend(platform=None):
--> 485   return _get_backend_uncached(platform)
    486 
    487 

/usr/local/lib/python3.9/dist-packages/jax/_src/xla_bridge.py in _get_backend_uncached(platform)
    473     if backend is None:
    474       if platform in _backends_errors:
--> 475         raise RuntimeError(f"Backend '{platform}' failed to initialize: "
    476                            f"{_backends_errors[platform]}")
    477       raise RuntimeError(f"Unknown backend {platform}")

RuntimeError: Backend 'tpu_driver' failed to initialize: Could not initialize backend 'tpu_driver'

我已确认TPU可用，怀疑问题出在代码行jax.tools.colab_tpu.setup_tpu('tpu_driver_20221011')，目前正在查阅Jax API寻求解决方案。备注：Jax与Jaxlib版本为0.4.6。

内容的提问来源于stack exchange，提问作者RanWang