$ docker run --rm rocm/jax:rocm7.0-jax0.6.0-py3.10 python -c 'import jax'
Traceback (most recent call last):
File "<string>", line 1, in <module>
ModuleNotFoundError: No module named 'jax'
$ docker run --rm rocm/jax:rocm7.0-jax0.6.0-py3.10 pip list
Package Version
-------------------- -----------
absl-py 2.3.1
attrs 25.3.0
auditwheel 6.4.2
beautifulsoup4 4.13.5
build 1.3.0
cloudpickle 3.1.1
contourpy 1.3.2
cycler 0.12.1
exceptiongroup 1.3.0
flatbuffers 25.2.10
fonttools 4.59.2
hypothesis 6.138.17
iniconfig 2.1.0
jax-rocm7-pjrt 0.6.0
jax-rocm7-plugin 0.6.0
Jinja2 3.1.6
kiwisolver 1.4.9
MarkupSafe 3.0.2
matplotlib 3.10.6
numpy 1.26.4
packaging 25.0
pillow 11.3.0
pip 25.2
pluggy 1.6.0
portpicker 1.6.0
psutil 7.0.0
pyelftools 0.32
Pygments 2.19.2
pyparsing 3.2.4
pyproject_hooks 1.2.0
pytest 8.4.2
pytest-html 4.1.1
pytest_html_merger 0.1.0
pytest-json-report 1.5.0
pytest-metadata 3.1.1
pytest-reportlog 0.4.0
pytest-rerunfailures 16.0.1
python-dateutil 2.9.0.post0
scipy 1.15.3
setuptools 80.9.0
six 1.17.0
sortedcontainers 2.4.0
soupsieve 2.8
tomli 2.2.1
typing_extensions 4.15.0
wheel 0.45.1
$ docker run --rm rocm/jax:rocm7.0-jax0.6.0-py3.12 python -c 'import jax'
Traceback (most recent call last):
File "<string>", line 1, in <module>
ModuleNotFoundError: No module named 'jax'
$ docker run --rm rocm/jax:rocm7.0-jax0.6.0-py3.12 pip list
Package Version
-------------------- -----------
absl-py 2.3.1
attrs 25.3.0
auditwheel 6.4.2
beautifulsoup4 4.13.5
build 1.3.0
cloudpickle 3.1.1
contourpy 1.3.3
cycler 0.12.1
flatbuffers 25.2.10
fonttools 4.59.2
hypothesis 6.138.17
iniconfig 2.1.0
jax-rocm7-pjrt 0.6.0
jax-rocm7-plugin 0.6.0
Jinja2 3.1.6
kiwisolver 1.4.9
MarkupSafe 3.0.2
matplotlib 3.10.6
numpy 1.26.4
packaging 25.0
pillow 11.3.0
pip 25.2
pluggy 1.6.0
portpicker 1.6.0
psutil 7.0.0
pyelftools 0.32
Pygments 2.19.2
pyparsing 3.2.4
pyproject_hooks 1.2.0
pytest 8.4.2
pytest-html 4.1.1
pytest_html_merger 0.1.0
pytest-json-report 1.5.0
pytest-metadata 3.1.1
pytest-reportlog 0.4.0
pytest-rerunfailures 16.0.1
python-dateutil 2.9.0.post0
scipy 1.16.2
setuptools 80.9.0
six 1.17.0
sortedcontainers 2.4.0
soupsieve 2.8
typing_extensions 4.15.0
wheel 0.45.1
Apologies if I'm misunderstanding how these containers are supposed to work or filing this issue in the incorrect location. I was surprised when trying to use the rocm/jax docker images that (at least) the following images don't have jax installed:
Details