Comments (9)
How did you install JAX?
from 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
from jax.
Can you show me your pip freeze
?
from jax.
$ 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
from jax.
Can you try running JAX_DEBUG_LOG_MODULES=jax._src.xla_bridge python -c 'import jax; print(jax.devices())'
and paste the output here?
from jax.
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 jax.
From the logs I realised the actual reason is that the TPU is used by another process. It works after the process is killed.
from jax.
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...
from jax.
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.
from jax.
Related Issues (20)
- Numerical precision differences in jitted vs. non-jitted code! HOT 3
- JAX indexing should support zero-dimensional boolean masks
- Memory profiler missing in trace HOT 10
- Getting some problems in the usage of jnp.linalg.inv HOT 3
- Installing jaxlib with CUDA Version: 11.7 HOT 5
- AssertionError raised due to failing assertion in partial_eval.py HOT 4
- Extending JAX with CUDA: cstdint: error: the global scope has no "..." HOT 3
- Complex gradient HOT 3
- Reconstructing a custom registered dataclass seems to give a segmentation fault
- AttributeError: module 'jax.random' has no attribute 'KeyArray' HOT 5
- Accessing incoming gradient and activations in custom_vjp HOT 4
- CUDA `XlaRuntimeError` with MPI on `jax==0.4.31` HOT 5
- Finite precision error calculations always 0 under JIT with bfloat16 HOT 3
- Errors from interaction between `shard_map` with `auto` and `lax.map`
- Document Named Collectives HOT 2
- local_device_count fails on wsl with 'cpu' HOT 4
- Scatter on Sharded Matrices has bugs HOT 1
- Port all custom calls to FFI
- custom_jvp leaks tracers if they're marked as nondiff_argnum HOT 2
- Add support nogil mode in JAX for Python 3.13
Recommend Projects
-
React
A declarative, efficient, and flexible JavaScript library for building user interfaces.
-
Vue.js
🖖 Vue.js is a progressive, incrementally-adoptable JavaScript framework for building UI on the web.
-
Typescript
TypeScript is a superset of JavaScript that compiles to clean JavaScript output.
-
TensorFlow
An Open Source Machine Learning Framework for Everyone
-
Django
The Web framework for perfectionists with deadlines.
-
Laravel
A PHP framework for web artisans
-
D3
Bring data to life with SVG, Canvas and HTML. 📊📈🎉
-
Recommend Topics
-
javascript
JavaScript (JS) is a lightweight interpreted programming language with first-class functions.
-
web
Some thing interesting about web. New door for the world.
-
server
A server is a program made to process requests and deliver data to clients.
-
Machine learning
Machine learning is a way of modeling and interpreting data that allows a piece of software to respond intelligently.
-
Visualization
Some thing interesting about visualization, use data art
-
Game
Some thing interesting about game, make everyone happy.
Recommend Org
-
Facebook
We are working to build community through open source technology. NB: members must have two-factor auth.
-
Microsoft
Open source projects and samples from Microsoft.
-
Google
Google ❤️ Open Source for everyone.
-
Alibaba
Alibaba Open Source for everyone
-
D3
Data-Driven Documents codes.
-
Tencent
China tencent open source team.
from jax.