run jax[tpu] without tpu

[cnsgsz]

I use cloud TPU VMs, but also run tests locally. As I try to use bazel rules_python's pip_install, which bundles python dependencies with the main script, I ran into issues running test locally as JAX cannot find TPU. Is there a way to turn off the tpu features in jax[tpu] so that I can run it without TPU?

What I tried:
I'm about to import jax and see CPU as the only device, but cannot run computation as JAX seems to keep trying to talk to TPU.
I read this issue which seem related.
Seem to me the only difference between jax[cpu] and jax[tpu] is two extra libs/packages.

[hawkinsp]

Can you say more about what goes wrong? What, exactly did you run? What happens when you try?

[cnsgsz]

Installed jax[tpu] locally

$ pip freeze
absl-py==1.2.0
apache-beam==2.40.0
astunparse==1.6.3
cached-property==1.5.2
cachetools==4.2.4
certifi==2022.6.15
charset-normalizer==2.1.0
chex==0.1.3
cloudpickle==2.1.0
clu==0.0.7
colorama==0.4.5
commonmark==0.9.1
contextlib2==21.6.0
crcmod==1.7
cycler==0.11.0
dill==0.3.1.1
dm-tree==0.1.7
docopt==0.6.2
etils==0.6.0
fastavro==1.5.3
fasteners==0.17.3
flatbuffers==1.12
flax==0.5.2
fonttools==4.34.4
gast==0.4.0
google-api-core==1.32.0
google-api-python-client==2.54.0
google-apitools==0.5.31
google-auth==1.35.0
google-auth-httplib2==0.1.0
google-auth-oauthlib==0.4.6
google-cloud-bigquery==2.34.4
google-cloud-bigquery-storage==2.14.1
google-cloud-bigtable==1.7.2
google-cloud-core==1.7.3
google-cloud-datastore==1.15.5
google-cloud-dlp==3.8.0
google-cloud-language==1.3.2
google-cloud-pubsub==2.13.4
google-cloud-pubsublite==1.4.2
google-cloud-recommendations-ai==0.2.0
google-cloud-spanner==1.19.3
google-cloud-videointelligence==1.16.3
google-cloud-vision==1.0.2
google-crc32c==1.3.0
google-pasta==0.2.0
google-resumable-media==2.3.3
googleapis-common-protos==1.56.4
grpc-google-iam-v1==0.12.4
grpcio==1.47.0
grpcio-gcp==0.2.2
grpcio-status==1.47.0
gviz-api==1.10.0
h5py==3.7.0
hdfs==2.7.0
httplib2==0.20.4
idna==3.3
importlib-metadata==4.12.0
importlib-resources==5.8.0
jax==0.3.14
jaxlib==0.3.14
keras==2.9.0
Keras-Preprocessing==1.1.2
kiwisolver==1.4.4
libclang==14.0.1
libtpu-nightly==0.1.dev20220627
Markdown==3.4.1
matplotlib==3.5.2
ml-collections==0.1.1
msgpack==1.0.4
numpy==1.22.4
oauth2client==4.1.3
oauthlib==3.2.0
opt-einsum==3.3.0
optax==0.1.3
orjson==3.7.8
overrides==6.1.0
packaging==21.3
Pillow==9.2.0
pkg_resources==0.0.0
promise==2.3
proto-plus==1.20.6
protobuf==3.19.4
pyarrow==7.0.0
pyasn1==0.4.8
pyasn1-modules==0.2.8
pydot==1.4.2
Pygments==2.12.0
pymongo==3.12.3
pyparsing==3.0.9
python-dateutil==2.8.2
pytz==2022.1
PyYAML==6.0
requests==2.28.1
requests-oauthlib==1.3.1
rich==11.2.0
rsa==4.9
scipy==1.8.1
six==1.16.0
tensorboard==2.9.1
tensorboard-data-server==0.6.1
tensorboard-plugin-profile==2.8.0
tensorboard-plugin-wit==1.8.1
tensorflow==2.9.1
tensorflow-datasets==4.6.0
tensorflow-estimator==2.9.0
tensorflow-io-gcs-filesystem==0.26.0
tensorflow-metadata==1.9.0
termcolor==1.1.0
toml==0.10.2
toolz==0.12.0
tqdm==4.64.0
typing-utils==0.1.0
typing_extensions==4.3.0
uritemplate==4.1.1
urllib3==1.26.10
Werkzeug==2.1.2
wrapt==1.14.1
zipp==3.8.1

Follow an example from github.com/google/jax

import jax has a warning, but fine otherwise.

>>> from jax import grad
/home/swang/venv/moffett_20220721tpu2/lib/python3.8/site-packages/jax/__init__.py:27: UserWarning: cloud_tpu_init failed: KeyError('')
 This a JAX bug; please report an issue at https://github.com/google/jax/issues
  _warn(f"cloud_tpu_init failed: {repr(exc)}\n This a JAX bug; please report "
>>> import jax.numpy as jnp
>>> 
>>> def tanh(x):  # Define a function
...   y = jnp.exp(-2.0 * x)
...   return (1.0 - y) / (1.0 + y)
... 
>>> grad_tanh = grad(tanh)  # Obtain its gradient function

Calling devices() for the first time has errors, but not subsequent ones

>>> jax.devices()
WARNING: Logging before InitGoogle() is written to STDERR
I0000 00:00:1658969750.942706    7367 common_lib.cc:144] Failed to fetch URL on try 1 out of 6: Couldn't connect to server
I0000 00:00:1658969754.014705    7367 common_lib.cc:144] Failed to fetch URL on try 2 out of 6: Couldn't connect to server
I0000 00:00:1658969757.086732    7367 common_lib.cc:144] Failed to fetch URL on try 3 out of 6: Couldn't connect to server
I0000 00:00:1658969760.158768    7367 common_lib.cc:144] Failed to fetch URL on try 4 out of 6: Couldn't connect to server
^C^CI0000 00:00:1658969763.230709    7367 common_lib.cc:144] Failed to fetch URL on try 5 out of 6: Couldn't connect to server
I0000 00:00:1658969766.302714    7367 common_lib.cc:144] Failed to fetch URL on try 6 out of 6: Couldn't connect to server
Failed to get accelerator type with error: Couldn't connect to serverI0000 00:00:1658969769.374712    7367 common_lib.cc:144] Failed to fetch URL on try 1 out of 6: Couldn't connect to server
I0000 00:00:1658969772.446723    7367 common_lib.cc:144] Failed to fetch URL on try 2 out of 6: Couldn't connect to server


>>> jax.devices()
[CpuDevice(id=0)]
>>> jax.devices("tpu")
Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/home/swang/venv/moffett_20220721tpu2/lib/python3.8/site-packages/jax/_src/lib/xla_bridge.py", line 449, in devices
    return get_backend(backend).devices()
  File "/home/swang/venv/moffett_20220721tpu2/lib/python3.8/site-packages/jax/_src/lib/xla_bridge.py", line 391, in get_backend
    return _get_backend_uncached(platform)
  File "/home/swang/venv/moffett_20220721tpu2/lib/python3.8/site-packages/jax/_src/lib/xla_bridge.py", line 383, in _get_backend_uncached
    raise RuntimeError(f"Unknown backend {platform}")
RuntimeError: Unknown backend tpu
>>> jax.devices("cpu")
[CpuDevice(id=0)]

Computation seems to work. If I haven't run jax.devices() earlier, I think this will show the same error running jax.devices for the first time showed.

>>> print(grad_tanh(1.0))   # Evaluate it at x = 1.0
0.4199743

[hawkinsp]

Thanks. Everything is working as intended: you are indeed able to run CPU computations.

All that needs to be fixed is we shouldn't print all the scary warning messages: if there's no TPU, we should immediately fail and do so politely. I filed #11668 for that.

Hope that helps!