NVIDIA
NVIDIA
JAX
Container
NVIDIA
NVIDIA
JAX

JAX is a framework for high-performance numerical computing and machine learning research. It includes Numpy-like APIs, automatic differentiation, XLA acceleration and simple primitives for scaling across GPUs and supports an ecosystem of libraries.

LayerLabelCreated
f912ae76a87bfcdf7d077708689938ceb14abeaf44482211d9a2e712ed6df113CONFIG
Entrypoint /opt/nvidia/nvidia_entrypoint.sh; WorkingDir /opt/maxtext
10/23/2024 6:37 AM UTC
310e1e4a004f4adb1141379a10fda2b2b397e9d976f582275318662f4201e884RUN
NVIDIA_BUILD_ID=117926700 URLREF_MAXTEXT=https://github.com/google/maxtext.git#20f57fb7138e43e35666e2037bf5cbc32043364a SRC_PATH_MAXTEXT=/opt/maxtext URLREF_JETSTREAM=https://github.com/google/jetstream.git#530d364e18d5759d7afd538bbb1d767de10a3f4c SRC_PATH_JETSTREAM=/opt/jetstream pip-finalize.sh
10/23/2024 6:37 AM UTC
c3e9b76edd5e8073833dd6c28afd596d71cce86341771069d46848cf59e5d3caRUN
RUN |5 NVIDIA_BUILD_ID=117926700 URLREF_MAXTEXT=https://github.com/google/maxtext.git#20f57fb7138e43e35666e2037bf5cbc32043364a SRC_PATH_MAXTEXT=/opt/maxtext URLREF_JETSTREAM=https://github.com/google/jetstream.git#530d364e18d5759d7afd538bbb1d767de10a3f4c SRC_PATH_JETSTREAM=/opt/jetstream /bin/sh -c <<"EOF" bash -ex git-clone.sh ${URLREF_JETSTREAM} ${SRC_PATH_JETSTREAM} echo "-e file://${SRC_PATH_JETSTREAM}" >> /opt/pip-tools.d/requirements-jetstream.in EOF # buildkit
10/23/2024 6:34 AM UTC
c4b33018b068556d7152bf9a5f16f13b08c0e03fa4c322e4c8099ba7501207c4ARG
SRC_PATH_JETSTREAM=/opt/jetstream
10/23/2024 6:34 AM UTC
2403a41cbadfb629133a1a83e43acd25a44c1aebfaec41246d767536c9e30aa9ARG
URLREF_JETSTREAM=https://github.com/google/jetstream.git#530d364e18d5759d7afd538bbb1d767de10a3f4c
10/23/2024 6:34 AM UTC
67a8f2d4c321ccc9fe7d30d32c9e203340231801fd49b785a891e9dbb855722eADD
./jax-toolbox/rosetta /opt/rosetta
10/23/2024 6:34 AM UTC
254929b834ce9f513f6f85cd0e2d8db0cc4dc4bd24c919859cc9e9032e5152eaADD
test-maxtext.sh /usr/local/bin
10/23/2024 6:34 AM UTC
574c412150fa5363f5670269f8c443735530075c60933c513c77094a972bb0f0RUN
RUN |3 NVIDIA_BUILD_ID=117926700 URLREF_MAXTEXT=https://github.com/google/maxtext.git#20f57fb7138e43e35666e2037bf5cbc32043364a SRC_PATH_MAXTEXT=/opt/maxtext /bin/sh -c <<"EOF" bash -ex git-clone.sh ${URLREF_MAXTEXT} ${SRC_PATH_MAXTEXT} sed -i '/google-jetstream/d' ${SRC_PATH_MAXTEXT}/requirements.txt echo "-r ${SRC_PATH_MAXTEXT}/requirements.txt" >> /opt/pip-tools.d/requirements-maxtext.in EOF # buildkit
10/23/2024 6:34 AM UTC
e9a5280fe48e80ea3c1b42486d583fd2b3ce631819ee195acc6d3e33cebca03cARG
SRC_PATH_MAXTEXT=/opt/maxtext
10/23/2024 6:34 AM UTC
eafe74df4330174235e3ed26abbf0887cd967848e4a96d9ff0917dc35bf3a961ARG
URLREF_MAXTEXT=https://github.com/google/maxtext.git#20f57fb7138e43e35666e2037bf5cbc32043364a
10/23/2024 6:34 AM UTC
...