TensorFlow can find GPU but JAX can't?

Sorry if this has already been asked, but I’m running into a problem where JAX can’t find my GPU:

venv/lib/python3.7/site-packages/jax/lib/xla_bridge.py:130: UserWarning: No GPU/TPU found, falling back to CPU.

However, TensorFlow has no problem finding the GPU. When I run tf.config.list_physical_devices('GPU'), I get the following output:

2020-12-28 20:31:33.309894: I tensorflow/compiler/xla/service/service.cc:168] XLA service 0x5646344172e0 initialized for platform CUDA (this does not guarantee that XLA will be used). Devices:
2020-12-28 20:31:33.309918: I tensorflow/compiler/xla/service/service.cc:176]   StreamExecutor device (0): Tesla V100-PCIE-32GB, Compute Capability 7.0
2020-12-28 20:31:33.311146: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1716] Found device 0 with properties:
pciBusID: 0000:d8:00.0 name: Tesla V100-PCIE-32GB computeCapability: 7.0
coreClock: 1.38GHz coreCount: 80 deviceMemorySize: 31.75GiB deviceMemoryBandwidth: 836.37GiB/s
2020-12-28 20:31:33.311171: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudart.so.10.1
2020-12-28 20:31:33.311201: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcublas.so.10
2020-12-28 20:31:33.311211: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcufft.so.10
2020-12-28 20:31:33.311220: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcurand.so.10
2020-12-28 20:31:33.311229: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcusolver.so.10
2020-12-28 20:31:33.311238: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcusparse.so.10
2020-12-28 20:31:33.311248: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudnn.so.7
2020-12-28 20:31:33.313354: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1858] Adding visible gpu devices: 0
2020-12-28 20:31:33.313382: I tensorflow/stream_executor/platform/default/dso_loader.cc:48] Successfully opened dynamic library libcudart.so.10.1
2020-12-28 20:31:33.850156: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1257] Device interconnect StreamExecutor with strength 1 edge matrix:
2020-12-28 20:31:33.850186: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1263]      0
2020-12-28 20:31:33.850192: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1276] 0:   N
2020-12-28 20:31:33.852686: I tensorflow/core/common_runtime/gpu/gpu_device.cc:1402] Created TensorFlow device (/job:localhost/replica:0/task:0/device:GPU:0 with 30132 MB memory) -> physical GPU (device: 0, name: Tesla V100-PCIE-32GB, pci bus id: 0000:d8:00.0, compute capability: 7.0)

[jakevdp]

I found that running

pip install pip==20.3.3
pip install --upgrade jax jaxlib==0.1.52+cuda101 -f https://storage.googleapis.com/jax-releases/jax_releases.html

reproduces the different version in metadata error, but using older versions of pip do not result in this error. For example, this installs correctly:

pip install pip==20.2.4
pip install --upgrade jax jaxlib==0.1.52+cuda101 -f https://storage.googleapis.com/jax-releases/jax_releases.html

I suspect that the pip 20.3 dependency resolver places stricter requirements on the contents of package metadata.

[jakevdp]

I don’t know of any workaround, but others might.