search

jax-ml / jax

Composable transformations of Python+NumPy programs: differentiate, vectorize, JIT to GPU/TPU, and more
http://jax.readthedocs.io/
Apache License 2.0
30.45k stars 2.8k forks source link

jaxlib or libtpu not detected on TPU Pod #22070

Open ayaka14732 opened 4 months ago

ayaka14732 commented 4 months ago

Description

I have installed the TPU version of JAX (including jaxlib and libtpu) on all hosts of a TPU Pod inside a venv. Then, I run the following command on all hosts:

. ~/venv/bin/activate; python -c 'import jax; print(jax.devices())'

I got this error:

A Google TPU may be present on this machine, but either a TPU-enabled jaxlib or libtpu is not installed. Falling back to cpu.

System info (python version, jaxlib version, accelerator, etc.)

>>> import jax; jax.print_environment_info()
A Google TPU may be present on this machine, but either a TPU-enabled jaxlib or libtpu is not installed. Falling back to cpu.
jax:    0.4.30
jaxlib: 0.4.30
numpy:  1.26.4
python: 3.12.4 (main, Jun  8 2024, 18:29:57) [GCC 11.4.0]
jax.devices (1 total, 1 local): [CpuDevice(id=0)]
process_count: 1
platform: uname_result(system='Linux', node='t1v-n-d95163c5-w-0', release='5.19.0-1022-gcp', version='#24~22.04.1-Ubuntu SMP Sun Apr 23 09:51:08 UTC 2023', machine='x86_64')
$ ls /dev/accel*
/dev/accel0  /dev/accel1  /dev/accel2  /dev/accel3
yashk2810 commented 4 months ago

How did you install JAX?

ayaka14732 commented 4 months ago

@yashk2810 I ran these commands on all hosts:

python3.12 -m venv ~/venv
. ~/venv/bin/activate
pip install -U pip
pip install -U wheel
pip install "jax[tpu]" -f https://storage.googleapis.com/jax-releases/libtpu_releases.html
yashk2810 commented 4 months ago

Can you show me your pip freeze?

ayaka14732 commented 4 months ago

@yashk2810 

$ pip freeze
certifi==2024.6.2
charset-normalizer==3.3.2
idna==3.7
jax==0.4.30
jaxlib==0.4.30
libtpu-nightly==0.1.dev20240617
ml-dtypes==0.4.0
numpy==2.0.0
opt-einsum==3.3.0
requests==2.32.3
scipy==1.14.0
urllib3==2.2.2
wheel==0.43.0
skye commented 4 months ago

Can you try running JAX_DEBUG_LOG_MODULES=jax._src.xla_bridge python -c 'import jax; print(jax.devices())' and paste the output here?

ayaka14732 commented 4 months ago

@skye 

DEBUG:2024-06-25 00:34:55,039:jax._src.xla_bridge:575: No jax_plugins namespace packages available
DEBUG:2024-06-25 00:34:55,049:jax._src.xla_bridge:969: Initializing backend 'cpu'
DEBUG:2024-06-25 00:34:55,109:jax._src.xla_bridge:981: Backend 'cpu' initialized
DEBUG:2024-06-25 00:34:55,109:jax._src.xla_bridge:969: Initializing backend 'cuda'
INFO:2024-06-25 00:34:55,109:jax._src.xla_bridge:889: Unable to initialize backend 'cuda': 
DEBUG:2024-06-25 00:34:55,109:jax._src.xla_bridge:969: Initializing backend 'rocm'
INFO:2024-06-25 00:34:55,109:jax._src.xla_bridge:889: Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
DEBUG:2024-06-25 00:34:55,109:jax._src.xla_bridge:969: Initializing backend 'tpu'
INFO:2024-06-25 00:34:55,142:jax._src.xla_bridge:889: Unable to initialize backend 'tpu': ABORTED: The TPU is already in use by process with pid 176028. Not attempting to load libtpu.so in this process.
WARNING:2024-06-25 00:34:55,143:jax._src.xla_bridge:940: A Google TPU may be present on this machine, but either a TPU-enabled jaxlib or libtpu is not installed. Falling back to cpu.
[CpuDevice(id=0)]
DEBUG:2024-06-25 00:34:55,812:jax._src.xla_bridge:575: No jax_plugins namespace packages available
DEBUG:2024-06-25 00:34:55,813:jax._src.xla_bridge:575: No jax_plugins namespace packages available
DEBUG:2024-06-25 00:34:55,836:jax._src.xla_bridge:575: No jax_plugins namespace packages available
DEBUG:2024-06-25 00:34:55,866:jax._src.xla_bridge:969: Initializing backend 'cpu'
DEBUG:2024-06-25 00:34:55,875:jax._src.xla_bridge:969: Initializing backend 'cpu'
DEBUG:2024-06-25 00:34:55,919:jax._src.xla_bridge:969: Initializing backend 'cpu'
DEBUG:2024-06-25 00:34:55,928:jax._src.xla_bridge:981: Backend 'cpu' initialized
DEBUG:2024-06-25 00:34:55,928:jax._src.xla_bridge:969: Initializing backend 'cuda'
INFO:2024-06-25 00:34:55,928:jax._src.xla_bridge:889: Unable to initialize backend 'cuda': 
DEBUG:2024-06-25 00:34:55,928:jax._src.xla_bridge:969: Initializing backend 'rocm'
INFO:2024-06-25 00:34:55,928:jax._src.xla_bridge:889: Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
DEBUG:2024-06-25 00:34:55,928:jax._src.xla_bridge:969: Initializing backend 'tpu'
DEBUG:2024-06-25 00:34:55,936:jax._src.xla_bridge:981: Backend 'cpu' initialized
DEBUG:2024-06-25 00:34:55,936:jax._src.xla_bridge:969: Initializing backend 'cuda'
INFO:2024-06-25 00:34:55,936:jax._src.xla_bridge:889: Unable to initialize backend 'cuda': 
DEBUG:2024-06-25 00:34:55,936:jax._src.xla_bridge:969: Initializing backend 'rocm'
INFO:2024-06-25 00:34:55,936:jax._src.xla_bridge:889: Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
DEBUG:2024-06-25 00:34:55,936:jax._src.xla_bridge:969: Initializing backend 'tpu'
INFO:2024-06-25 00:34:55,960:jax._src.xla_bridge:889: Unable to initialize backend 'tpu': ABORTED: The TPU is already in use by process with pid 463521. Not attempting to load libtpu.so in this process.
WARNING:2024-06-25 00:34:55,961:jax._src.xla_bridge:940: A Google TPU may be present on this machine, but either a TPU-enabled jaxlib or libtpu is not installed. Falling back to cpu.
[CpuDevice(id=0)]
INFO:2024-06-25 00:34:55,970:jax._src.xla_bridge:889: Unable to initialize backend 'tpu': ABORTED: The TPU is already in use by process with pid 482056. Not attempting to load libtpu.so in this process.
WARNING:2024-06-25 00:34:55,971:jax._src.xla_bridge:940: A Google TPU may be present on this machine, but either a TPU-enabled jaxlib or libtpu is not installed. Falling back to cpu.
[CpuDevice(id=0)]
DEBUG:2024-06-25 00:34:55,982:jax._src.xla_bridge:981: Backend 'cpu' initialized
DEBUG:2024-06-25 00:34:55,982:jax._src.xla_bridge:969: Initializing backend 'cuda'
INFO:2024-06-25 00:34:55,982:jax._src.xla_bridge:889: Unable to initialize backend 'cuda': 
DEBUG:2024-06-25 00:34:55,982:jax._src.xla_bridge:969: Initializing backend 'rocm'
INFO:2024-06-25 00:34:55,982:jax._src.xla_bridge:889: Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
DEBUG:2024-06-25 00:34:55,982:jax._src.xla_bridge:969: Initializing backend 'tpu'
INFO:2024-06-25 00:34:56,015:jax._src.xla_bridge:889: Unable to initialize backend 'tpu': ABORTED: The TPU is already in use by process with pid 149052. Not attempting to load libtpu.so in this process.
WARNING:2024-06-25 00:34:56,016:jax._src.xla_bridge:940: A Google TPU may be present on this machine, but either a TPU-enabled jaxlib or libtpu is not installed. Falling back to cpu.
[CpuDevice(id=0)]
ayaka14732 commented 4 months ago

From the logs I realised the actual reason is that the TPU is used by another process. It works after the process is killed.

skye commented 4 months ago

Ah. This is supposed to be raised as an exception instead of falling back to CPU. That functionality must have regressed. Now to figure out why...

s-smits commented 4 months ago 
air@MacBook-Air-van-Air ~ % pip freeze
absl-py==2.0.0
accelerate==0.31.0
addict==2.4.0
aenum==3.1.15
aifs==0.0.9
aiofiles==23.2.1
aiohttp==3.9.3
aiosignal==1.3.1
albumentations==1.3.1
alembic==1.13.1
altair==5.3.0
annotated-types==0.6.0
ansi2html==1.8.0
anthropic==0.18.1
antlr4-python3-runtime==4.9.3
anyascii==0.3.2
anyio==3.7.1
appdirs==1.4.4
appnope==0.1.3
argon2-cffi==23.1.0
argon2-cffi-bindings==21.2.0
arrow==1.3.0
asgiref==3.7.2
asteroid-filterbanks==0.4.0
astor==0.8.1
asttokens==2.4.1
astunparse==1.6.3
async-lru==2.0.4
async-timeout==4.0.3
attrs==23.2.0
audioread==3.0.1
av==12.0.0
Babel==2.14.0
backcall==0.2.0
backoff==2.2.1
bangla==0.0.2
basicsr==1.4.2
bcrypt==4.1.2
beautifulsoup4==4.12.3
bitarray==2.8.2
bleach==6.1.0
blendmodes==2022
blessed==1.20.0
blinker==1.8.2
blis==0.7.11
bnnumerizer==0.0.2
bnunicodenormalizer==0.1.6
boto3==1.28.80
botocore==1.31.80
branca==0.7.0
Brotli==1.1.0
build==1.0.3
CacheControl==0.13.1
cachetools==5.3.3
camelot-py==0.11.0
Cartopy==0.22.0
catalogue==2.0.10
certifi==2024.6.2
cffi @ file:///private/tmp/cffi-20231014-5454-1xfyd4s/cffi-1.16.0
cftime==1.6.3
chardet==5.2.0
charset-normalizer==3.3.2
chex==0.1.85
chroma-hnswlib==0.7.3
chromadb==0.4.24
ci-info==0.3.0
clean-fid==0.1.35
cleo==2.1.0
click==8.1.7
click-plugins==1.1.1
cligj==0.7.2
clip @ https://github.com/openai/CLIP/archive/d50d76daa670286dd6cacf3bcd80b5e4823fc8e1.zip#sha256=b5842c25da441d6c581b53a5c60e0c2127ebafe0f746f8e15561a006c6c3be6a
cloudpathlib==0.16.0
colorama==0.4.6
coloredlogs==15.0.1
colorlog==6.8.2
comm==0.2.0
confection==0.1.4
configobj==5.0.8
configparser==6.0.1
contourpy==1.2.0
coqpit==0.0.17
crashtest==0.4.1
crewai==0.22.5
crewai-tools==0.1.7
cryptography==42.0.2
cssselect2==0.7.0
ctranslate2==4.2.1
cycler==0.12.1
cymem==2.0.8
Cython==3.0.5
dataclasses-json==0.6.4
dataclasses-json-speakeasy==0.5.11
datasets==2.14.6
dateparser==1.1.8
debugpy==1.8.0
decorator==5.1.1
defusedxml==0.7.1
demoji==1.1.0
Deprecated==1.2.14
deprecation==2.1.0
depth_anything @ https://github.com/huchenlei/Depth-Anything/releases/download/v1.0.0/depth_anything-2024.1.22.0-py2.py3-none-any.whl#sha256=26c1d38b8c3c306b4a2197d725a4b989ff65f7ebcf4fb5a96a1b6db7fbd56780
diagrams==0.23.4
diffusers==0.27.2
dill==0.3.7
diskcache==5.6.3
distlib==0.3.8
distro==1.9.0
dm-tree==0.1.8
dnspython==2.6.1
docker==7.1.0
docopt==0.6.2
docstring-parser==0.15
dsine @ https://github.com/sdbds/DSINE/releases/download/1.0.2/dsine-2024.3.23-py3-none-any.whl#sha256=b9ea3bacce09f9b3f7fb4fa12471da7e465b2f9a60412711105a9238db280442
dulwich==0.21.7
dunamai==1.19.2
easydict==1.11
editor==1.6.5
effdet==0.4.1
einops==0.8.0
email_validator==2.1.1
embedchain==0.1.102
emoji==2.8.0
en-core-web-sm @ https://github.com/explosion/spacy-models/releases/download/en_core_web_sm-3.7.1/en_core_web_sm-3.7.1-py3-none-any.whl#sha256=86cc141f63942d4b2c5fcee06630fd6f904788d2f0ab005cce45aadb8fb73889
encodec==0.1.1
et-xmlfile==1.1.0
etelemetry==0.3.1
etils==1.6.0
evaluate==0.4.1
exceptiongroup==1.2.0
executing==2.0.1
fabric @ file:///Users/air/Repositories/fabric
facexlib==0.3.0
fairseq @ git+https://github.com/pytorch/fairseq.git@da8fb630880d529ab47e53381c30ddc8ad235216
fastapi==0.110.3
fastapi-cli==0.0.3
fastBPE==0.1.0
faster-whisper==1.0.2
fastjsonschema==2.19.0
ffmpy==0.3.2
filelock==3.14.0
filetype==1.2.0
filterpy==1.4.5
fiona==1.9.5
fitz==0.0.1.dev2
FLAML==2.1.2
Flask==3.0.3
flatbuffers==23.5.26
flax==0.7.5
folium==0.15.0
fonttools==4.44.0
fqdn==1.5.1
frozenlist==1.4.0
fsspec==2024.6.0
ftfy==6.2.0
funcy==2.0
future==1.0.0
fvcore==0.1.5.post20221221
g2pkk==0.1.2
gast==0.5.4
geffnet==1.0.2
gensim==4.3.2
geographiclib==2.0
geojson==2.5.0
geopandas==0.14.3
geopy==2.4.0
gfpgan==1.3.8
git-python==1.0.3
gitdb==4.0.11
GitPython==3.1.43
glob2==0.5
google-ai-generativelanguage==0.6.4
google-api-core==2.19.1
google-api-python-client==2.136.0
google-auth==2.31.0
google-auth-httplib2==0.2.0
google-auth-oauthlib==1.1.0
google-cloud-aiplatform==1.58.0
google-cloud-bigquery==3.25.0
google-cloud-core==2.4.1
google-cloud-resource-manager==1.12.3
google-cloud-storage==2.17.0
google-crc32c==1.5.0
google-generativeai==0.5.4
google-pasta==0.2.0
google-resumable-media==2.7.1
googleapis-common-protos==1.63.0
googlemaps==4.10.0
gptcache==0.1.43
gradio==4.26.0
gradio_client==0.15.1
graphviz==0.20.1
grpc-google-iam-v1==0.13.1
grpcio==1.64.1
grpcio-status==1.62.2
gruut==2.2.3
gruut-ipa==0.13.0
gruut-lang-de==2.0.0
gruut-lang-en==2.0.0
gruut-lang-es==2.0.0
gruut-lang-fr==2.0.2
h11==0.14.0
h5py==3.10.0
handrefinerportable @ https://github.com/huchenlei/HandRefinerPortable/releases/download/v1.0.1/handrefinerportable-2024.2.12.0-py2.py3-none-any.whl#sha256=1e6c702905919f4c49bcb2db7b20d334e8458a7555cd57630600584ec38ca6a9
hangul-romanize==0.1.0
helpers==0.2.0
html2image==2.0.4.3
html2text==2024.2.26
httpcore==1.0.5
httplib2==0.22.0
httptools==0.6.1
httpx==0.27.0
huggingface-hub==0.23.3
humanfriendly==10.0
hydra-core==1.3.2
HyperPyYAML==1.2.2
idna==3.7
imageio==2.33.1
importlib-metadata==6.11.0
importlib-resources==6.1.1
inflect==7.0.0
inflection==0.5.1
iniconfig==2.0.0
inquirer==3.2.1
insightface==0.7.3
install==1.3.5
installer==0.7.0
instructor==0.5.2
iopath==0.1.10
ip-adapter==0.1.0
ipykernel==6.27.1
ipython==8.12.3
ipython-genutils==0.2.0
ipywidgets==8.1.2
isodate==0.6.1
isoduration==20.11.0
itsdangerous==2.1.2
jamo==0.4.1
jaraco.classes==3.3.0
jax==0.4.28
jaxlib==0.4.28
jedi==0.19.1
jellyfish==0.11.2
jieba==0.42.1
Jinja2==3.1.4
jinja2-ansible-filters==1.3.2
jiter==0.4.2
jmespath==1.0.1
joblib==1.3.2
json5==0.9.25
jsonlines==1.2.0
jsonmerge==1.8.0
jsonpatch==1.33
jsonpath-python==1.0.6
jsonpointer==2.4
jsonschema==4.22.0
jsonschema-specifications==2023.12.1
julius==0.2.7
jupyter-events==0.9.0
jupyter-lsp==2.2.5
jupyter_client==8.6.0
jupyter_core==5.5.0
jupyter_server==2.12.5
jupyter_server_terminals==0.5.2
jupyterlab==4.2.1
jupyterlab_pygments==0.3.0
jupyterlab_server==2.27.2
jupyterlab_widgets==3.0.10
jwt==1.3.1
keras==2.15.0
keyring==24.3.0
kiwisolver==1.4.5
kornia==0.6.7
kubernetes==29.0.0
lancedb==0.5.7
langchain==0.1.11
langchain-community==0.0.25
langchain-core==0.1.31
langchain-openai==0.0.5
langchain-text-splitters==0.0.1
langcodes==3.3.0
langdetect==1.0.9
langsmith==0.1.23
lark==1.1.2
layoutparser==0.3.4
lazy_loader==0.3
libclang==16.0.6
librosa==0.10.0
lightning==2.2.4
lightning-utilities==0.11.2
linkify-it-py==2.0.3
litellm==1.31.3
llama_cpp_python==0.2.76
llvmlite==0.41.1
lmdb==1.4.1
looseversion==1.3.0
lxml==4.9.3
Mako==1.3.2
mapbox-earcut==1.0.1
Markdown==3.5.1
markdown-it-py==3.0.0
MarkupSafe==2.1.5
marshmallow==3.20.2
matplotlib==3.8.2
matplotlib-inline==0.1.6
mdit-py-plugins==0.3.3
mdurl==0.1.2
mediapipe==0.10.11
memory-profiler==0.61.0
meteostat==1.6.7
mistune==3.0.2
ml-dtypes==0.2.0
mlx==0.15.0
mlx-lm==0.14.3
mlx-vlm==0.0.5
mmh3==4.1.0
monotonic==1.6
more-itertools==10.2.0
MouseInfo==0.1.3
mpmath==1.3.0
msgpack==1.0.7
multidict==6.0.4
multiprocess==0.70.15
murmurhash==1.0.10
mutagen==1.47.0
mypy-extensions==1.0.0
namex==0.0.7
natsort==8.4.0
nbclassic==1.0.0
nbclient==0.9.0
nbconvert==7.12.0
nbformat==5.9.2
nest-asyncio==1.5.8
netCDF4==1.6.5
networkx==2.8.8
nibabel==5.2.1
nipype==1.8.6
nltk==3.8.1
nodeenv==1.9.1
notebook==7.2.0
notebook_shim==0.2.4
num2words==0.5.13
numba==0.58.1
numpy==1.26.4
oauthlib==3.2.2
ollama==0.1.9
omegaconf==2.2.3
onnx==1.15.0
onnxruntime==1.15.1
open-clip-torch==2.20.0
open-interpreter==0.2.2
openai==1.13.3
opencv-contrib-python==4.9.0.80
opencv-python==4.9.0.80
opencv-python-headless==4.9.0.80
openpyxl==3.1.2
opentelemetry-api==1.23.0
opentelemetry-exporter-otlp-proto-common==1.23.0
opentelemetry-exporter-otlp-proto-grpc==1.23.0
opentelemetry-exporter-otlp-proto-http==1.23.0
opentelemetry-instrumentation==0.44b0
opentelemetry-instrumentation-asgi==0.44b0
opentelemetry-instrumentation-fastapi==0.44b0
opentelemetry-proto==1.23.0
opentelemetry-sdk==1.23.0
opentelemetry-semantic-conventions==0.44b0
opentelemetry-util-http==0.44b0
opt-einsum==3.3.0
optax==0.1.7
optuna==3.5.0
orbax-checkpoint==0.4.8
orjson==3.9.15
outcome==1.3.0.post0
overrides==7.7.0
packaging==23.2
pandas==2.2.2
pandocfilters==1.5.0
parso==0.8.3
pathlib==1.0.1
pathspec==0.12.1
patsy==0.5.6
pdf2image==1.17.0
pdfdocument==4.0.0
pdfminer.six==20231228
pdfplumber==0.11.1
peft==0.10.0
pexpect==4.9.0
pickleshare==0.7.5
piexif==1.1.3
pikepdf==8.12.0
pillow==10.3.0
pip-tools==7.4.1
pipreqs==0.5.0
pkginfo==1.9.6
platformdirs==4.2.0
plotly==5.19.0
pluggy==1.5.0
plumbum==1.8.2
poetry==1.7.1
poetry-core==1.8.1
poetry-plugin-export==1.6.0
polars==0.20.2
pooch==1.8.0
portalocker==2.8.2
posthog==3.3.1
praisonAI==0.0.18
preshed==3.0.9
prettytable==3.9.0
primePy==1.3
prometheus_client==0.20.0
prompt-toolkit==3.0.36
proto-plus==1.24.0
protobuf==4.25.3
prov==2.0.0
psutil==5.9.5
ptyprocess==0.7.0
pulsar-client==3.4.0
pure-eval==0.2.2
py==1.11.0
pyannote.audio==3.2.0
pyannote.core==5.0.0
pyannote.database==5.1.0
pyannote.metrics==3.2.1
pyannote.pipeline==3.0.1
pyarrow==16.1.0
pyasn1==0.5.1
pyasn1-modules==0.3.0
pyautogen==0.2.31
PyAutoGUI==0.9.54
pycocotools==2.0.7
pycollada==0.8
pycparser==2.22
pycryptodomex==3.20.0
pydantic==2.7.1
pydantic-settings==2.2.1
pydantic_core==2.18.2
pydeck==0.9.1
pydeps==1.12.20
pydot==1.4.2
pydub==0.25.1
PyGetWindow==0.0.9
PyGithub==1.59.1
Pygments==2.18.0
PyJWT==2.8.0
pylance==0.9.18
PyMsgBox==1.0.9
PyMuPDF==1.24.7
PyMuPDFb==1.24.6
PyNaCl==1.5.0
pynndescent==0.5.11
pyobjc-core==10.0
pyobjc-framework-Cocoa==10.0
pyobjc-framework-Quartz==10.0
pyowm==3.3.0
pyparsing==3.1.1
pypdf==3.17.4
PyPDF2==3.0.1
pypdfium2==4.26.0
pyperclip==1.8.2
PyPika==0.48.9
pypinyin==0.50.0
pyproj==3.6.1
pyproject_hooks==1.0.0
PyRect==0.2.0
pyright==1.1.370
pysbd==0.3.4
PyScreeze==0.1.30
pyshp==2.3.1
PySocks==1.7.1
pytesseract==0.3.10
pytest==8.2.2
python-crfsuite==0.9.10
python-dateutil==2.9.0.post0
python-docx==1.1.0
python-dotenv==1.0.0
python-iso639==2024.1.2
python-json-logger==2.0.7
python-magic==0.4.27
python-multipart==0.0.9
python-pptx==0.6.23
pytils==0.4.1
pytorch-lightning==2.2.0
pytorch-metric-learning==2.5.0
pytorch-transformers==1.2.0
pytube==15.0.0
pytweening==1.0.7
pytz==2024.1
PyWavelets==1.6.0
pyxnat==1.6.2
PyYAML==6.0.1
pyyaml-include==1.3.2
pyzmq==25.1.1
qudida==0.0.4
questionary==2.0.1
rapidfuzz==3.6.1
ratelimiter==1.2.0.post0
rdflib==7.0.0
readchar==4.0.5
realesrgan==0.3.0
referencing==0.35.1
regex==2023.12.25
reportlab==4.1.0
requests==2.32.3
requests-oauthlib==1.3.1
requests-toolbelt==1.0.0
resize-right==0.0.2
responses==0.18.0
retry==0.9.2
rfc3339-validator==0.1.4
rfc3986-validator==0.1.1
rich==13.7.1
rpds-py==0.18.1
rsa==4.9
Rtree==1.2.0
ruamel.yaml==0.18.6
ruamel.yaml.clib==0.2.8
rubicon-objc==0.4.7
ruff==0.5.0
runs==1.2.0
s3transfer==0.7.0
sacrebleu==2.3.2
sacremoses==0.1.1
safetensors==0.4.3
schema==0.7.7
scikit-image==0.21.0
scikit-learn==1.3.2
scipy==1.11.4
seaborn==0.13.0
selenium==4.22.0
semantic-version==2.10.0
semver==3.0.2
Send2Trash==1.8.2
sentence-transformers==2.2.2
sentencepiece==0.2.0
setfit==0.7.0
shapely==2.0.2
shellingham==1.5.4
shtab==1.7.1
simplejson==3.19.2
six==1.16.0
smart-open==6.4.0
smmap==5.0.1
sniffio==1.3.0
sortedcontainers==2.4.0
sounddevice==0.4.6
soundfile==0.12.1
soupsieve==2.5
soxr==0.3.7
spacy==3.7.2
spacy-legacy==3.0.12
spacy-loggers==1.0.5
spandrel==0.1.6
speechbrain==1.0.0
SQLAlchemy==2.0.27
srsly==2.4.8
sse-starlette==2.1.0
stack-data==0.6.3
stanza==1.3.0
starlette==0.37.2
starlette-context==0.3.6
statsmodels==0.14.1
stdlib-list==0.10.0
stqdm==0.0.5
streamlit==1.35.0
SudachiDict-core==20240109
SudachiPy==0.6.8
svg.path==6.3
svglib==1.5.1
sympy==1.12
tabulate==0.9.0
tb-nightly==2.17.0a20240416
tenacity==8.3.0
tensorboard==2.15.1
tensorboard-data-server==0.7.2
tensorboardX==2.6.2.2
tensorflow==2.15.0
tensorflow-estimator==2.15.0
tensorflow-io-gcs-filesystem==0.34.0
tensorflow-macos==2.15.0
tensorstore==0.1.52
termcolor==1.1.0
terminado==0.18.0
text-hammer==0.1.5
textblob==0.17.1
tf-keras==2.15.0
thinc==8.2.1
threadpoolctl==3.2.0
throttle==0.2.2
tifffile==2023.12.9
tiktoken==0.5.2
timm==0.9.2
tinycss2==1.2.1
tokenizers==0.19.1
tokentrim==0.1.13
tomesd==0.1.3
toml==0.10.2
tomli==2.0.1
tomlkit==0.12.0
tools==0.1.9
toolz==0.12.1
torch==2.2.0
torch-audiomentations==0.11.1
torch-pitch-shift==1.2.4
torch-tb-profiler==0.4.3
torchaudio==2.2.0
torchdiffeq==0.2.3
torchmetrics==1.3.2
torchsde==0.2.6
torchvision==0.17.0
tornado==6.4.1
tqdm==4.66.4
trainer==0.0.36
traitlets==5.14.0
traits==6.3.2
trampoline==0.1.2
transformers==4.41.2
trimesh==3.21.7
trio==0.25.1
trio-websocket==0.11.1
trl==0.9.4
trove-classifiers==2024.1.31
TTS==0.22.0
twokenize==1.0.0
typed-ast==1.5.5
typer==0.9.4
types-python-dateutil==2.8.19.20240106
typing-inspect==0.9.0
typing_extensions==4.12.2
tyro==0.7.3
tzdata==2024.1
tzlocal==5.2
uc-micro-py==1.0.3
ujson==5.9.0
umap-learn==0.5.5
Unidecode==1.3.8
unstructured==0.10.25
unstructured-client==0.17.0
unstructured-inference==0.7.23
unstructured.pytesseract==0.3.12
uri-template==1.3.0
uritemplate==4.1.1
urllib3==2.2.1
uvicorn==0.27.1
uvloop==0.19.0
virtualenv==20.25.0
wasabi==1.1.2
watchdog==4.0.1
watchfiles==0.21.0
wcwidth==0.2.12
weasel==0.3.4
webcolors==1.13
webencodings==0.5.1
websocket-client==1.8.0
websockets==11.0.3
Werkzeug==3.0.1
wget==3.2
widgetsnbextension==4.0.10
wordcloud==1.9.3
wordsegment==1.3.1
wrapt==1.14.1
wsproto==1.2.0
xarray==2024.2.0
xattr==0.10.1
xgboost==2.0.3
xlrd==2.0.1
XlsxWriter==3.2.0
xlwt==1.3.0
xmltodict==0.13.0
xmod==1.8.1
xxhash==3.4.1
yacs==0.1.8
yapf==0.40.2
yarg==0.1.9
yarl==1.9.2
youtube-transcript-api==0.6.2
yt-dlp==2023.11.16
zipp==3.17.0
$ JAX_DEBUG_LOG_MODULES=jax._src.xla_bridge python -c 'import jax; print(jax.devices())'
DEBUG:2024-07-10 12:20:02,539:jax._src.xla_bridge:575: No jax_plugins namespace packages available
DEBUG:2024-07-10 12:20:02,546:jax._src.xla_bridge:969: Initializing backend 'cpu'
DEBUG:2024-07-10 12:20:02,609:jax._src.xla_bridge:981: Backend 'cpu' initialized
DEBUG:2024-07-10 12:20:02,609:jax._src.xla_bridge:969: Initializing backend 'cuda'
INFO:2024-07-10 12:20:02,609:jax._src.xla_bridge:889: Unable to initialize backend 'cuda':
DEBUG:2024-07-10 12:20:02,609:jax._src.xla_bridge:969: Initializing backend 'rocm'
INFO:2024-07-10 12:20:02,609:jax._src.xla_bridge:889: Unable to initialize backend 'rocm': module 'jaxlib.xla_extension' has no attribute 'GpuAllocatorConfig'
DEBUG:2024-07-10 12:20:02,609:jax._src.xla_bridge:969: Initializing backend 'tpu'
INFO:2024-07-10 12:20:02,633:jax._src.xla_bridge:889: Unable to initialize backend 'tpu': ABORTED: The TPU is already in use by another process probably owned by another user. Run "$ sudo lsof -w /dev/accel0" to figure out which process is using the TPU. If you still get this message, run "$ sudo rm /tmp/libtpu_lockfile".
WARNING:2024-07-10 12:20:02,634:jax._src.xla_bridge:940: A Google TPU may be present on this machine, but either a TPU-enabled jaxlib or libtpu is not installed. Falling back to cpu.