Open ayaka14732 opened 5 months ago
How did you install JAX?
@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
Can you show me your pip freeze
?
@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
Can you try running JAX_DEBUG_LOG_MODULES=jax._src.xla_bridge python -c 'import jax; print(jax.devices())'
and paste the output here?
@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)]
From the logs I realised the actual reason is that the TPU is used by another process. It works after the process is killed.
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...
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.
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:
I got this error:
System info (python version, jaxlib version, accelerator, etc.)