Stuck on an issue?

Lightrun Answers was designed to reduce the constant googling that comes with debugging 3rd party libraries. It collects links to all the places you might be looking at while hunting down a tough bug.

And, if you’re still stuck at the end, we’re happy to hop on a call to see how we can help out.

TPU not working

See original GitHub issue

Description

I created a new TPU VM and sucessfully used it but after stopping / starting the VM am getting the following error:

Python 3.8.10 (default, Jun 22 2022, 20:18:18) 
[GCC 9.4.0] on linux
Type "help", "copyright", "credits" or "license" for more information.
>>> import jax
>>> jax.devices()
Traceback (most recent call last):
  File "/home/chris/.local/lib/python3.8/site-packages/jax/_src/lib/xla_bridge.py", line 333, in backends
    backend = _init_backend(platform)
  File "/home/chris/.local/lib/python3.8/site-packages/jax/_src/lib/xla_bridge.py", line 385, in _init_backend
    backend = factory()
  File "/home/chris/.local/lib/python3.8/site-packages/jax/_src/lib/xla_bridge.py", line 191, in tpu_client_timer_callback
    client = xla_client.make_tpu_client()
  File "/home/chris/.local/lib/python3.8/site-packages/jaxlib/xla_client.py", line 122, in make_tpu_client
    return _xla.get_tpu_client(
jaxlib.xla_extension.XlaRuntimeError: NOT_FOUND: No ba16c7433 device found.

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/chris/.local/lib/python3.8/site-packages/jax/_src/lib/xla_bridge.py", line 483, in devices
    return get_backend(backend).devices()
  File "/home/chris/.local/lib/python3.8/site-packages/jax/_src/lib/xla_bridge.py", line 425, in get_backend
    return _get_backend_uncached(platform)
  File "/home/chris/.local/lib/python3.8/site-packages/jax/_src/lib/xla_bridge.py", line 409, in _get_backend_uncached
    bs = backends()
  File "/home/chris/.local/lib/python3.8/site-packages/jax/_src/lib/xla_bridge.py", line 350, in backends
    raise RuntimeError(err_msg)
RuntimeError: Unable to initialize backend 'tpu': NOT_FOUND: No ba16c7433 device found. (set JAX_PLATFORMS='' to automatically choose an available backend)

What jax/jaxlib version are you using?

jax 0.3.25[tpu], jaxlib 0.3.25

Which accelerator(s) are you using?

TPU

Additional system info

No response

NVIDIA GPU info

No response

Issue Analytics

State:
Created 10 months ago
Comments:9

Top GitHub Comments

1reaction

chrisfleshercommented, Dec 12, 2022

Yes, creating a new VM seems to work until we stop / start the instance. We’re using this as a workaround until this issue can be resolved.

0reactions

skyecommented, Nov 18, 2022

Thanks. We’re still looking into this. In the meantime and if you’re able, creating a new VM will likely resolve the issue (definitely let us know if it doesn’t).

Top Results From Across the Web

Troubleshooting Cloud TPU errors and performance issues

These troubleshooting documents describe error conditions and performance issues you might see while training with Cloud TPUs using TensorFlow, JAX, and ...

TPU not working in kernel | Data Science and Machine Learning

For a week now I cant get TPU to work in a kernel, it is basically not available. When I try to install...

TPU: Tips and Tricks

When this happens, the TPU polymeric chains break down as the material heats up, making it foamy and not smooth. This can cause...

3D Printing Troubleshooting Guide - NinjaTek

Too little infill: Adjust your minimum infill. Flexible filaments often require more infill than their more rigid counterparts. For NinjaFlex ®, try to...

TPU filament not loading - Prusa Forum

Re: TPU filament not loading ... Filament loading is far to fast for TPU and it gets stuck at the toothed gear. Use...