[Issue]: rocm/jax images don't have jax or jaxlib installed

[GMNGeoffrey]

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:

• rocm/jax:rocm7.0-jax0.6.0-py3.10
• rocm/jax:rocm7.0-jax0.6.0-py3.12

Details

$ 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

[GMNGeoffrey]

Ping on this? @charleshofer can you PTAL