rocm/jax:rocm10.1.0_ubuntu24.04_py3.12_jax_release_0.10.2

Manifest digest

sha256:d1690b21415461eac0b59687213a1d62d961691557d595eeb2e9137c564170bb

OS/ARCH

linux/amd64

Compressed size

7.68 GB

Last pushed

about 1 month by rocm

Type

Image

Manifest digest

sha256:d1690b21415461eac0b59687213a1d62d961691557d595eeb2e9137c564170bb

Image Layers

1ARG RELEASE0 B

2ARG LAUNCHPAD_BUILD_ARCH0 B

3LABEL org.opencontainers.image.version=24.040 B

4ADD file ... in / 29.21 MB

5CMD ["/bin/bash"]0 B

6ARG VERSION=10.1.00 B

7ARG JAX_INDEX_URL=https://d1kq5rufa68c0p.cloudfront.net/rocm/whl-next/0 B

8ARG ROCM_SPECS=rocm[libraries,device-all]==10.1.00 B

9ARG JAX_SPECS=jax==0.10.20 B

10ARG JAX_PLUGIN_SPECS=jax-rocm10-plugin==0.10.2+rocm10.1.0 jax-rocm10-pjrt==0.10.2+rocm10.1.00 B

11LABEL org.opencontainers.image.title=ROCm JAX (rocm-ai-docker-stack) org.opencontainers.image.description=Multi-arch0 B

12COPY staging/install_rocm_deps.sh /tmp/install_rocm_deps.sh # buildkit1.31 KB

13COPY scripts/smoke_jax.sh /usr/local/bin/smoke_jax.sh # buildkit747 B

14

RUN |5 VERSION=10.1.0 JAX_INDEX_URL=https://d1kq5rufa68c0p.cloudfront.net/rocm/whl-next/ ROCM_SPECS=rocm[libraries,device-all]==10.1.0

134.14 MB

15

RUN |5 VERSION=10.1.0 JAX_INDEX_URL=https://d1kq5rufa68c0p.cloudfront.net/rocm/whl-next/ ROCM_SPECS=rocm[libraries,device-all]==10.1.0

198 B

16

RUN |5 VERSION=10.1.0 JAX_INDEX_URL=https://d1kq5rufa68c0p.cloudfront.net/rocm/whl-next/ ROCM_SPECS=rocm[libraries,device-all]==10.1.0

7.52 GB

17COPY --chmod=0755 scripts/setup_jax_env.sh /usr/local/bin/setup_jax_env #9.74 KB

18ENV PYTHONUNBUFFERED=10 B

Command

 ARG RELEASE