ContainersNVIDIADeepSeek-R1NVIDIALlama-3.1-Nemotron-70B-InstructNVIDIAPyTorch
CollectionsNVIDIAOmniverse Kit (FB)NVIDIADeepStream SDKNVIDIAOmniverse Kit App Streaming
ModelsNVIDIAStyleGAN3 pretrained modelsNVIDIAPeopleNetNVIDIATrafficCamNet
ResourcesNVIDIARiva Skills Quick Start
Helm ChartsNVIDIAGPU OperatorNVIDIANVIDIA NIM Operator
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
sha256:56369f2a9bd40e47494aaa92570a765e7e1de820e840f194c66a8bad23125acaRUN
pip-finalize.sh
05/08/2024 5:03 AM UTC
sha256:4f4fb700ef54461cfa02571ae0db9a0dc1e0cdb5577484a6d75e68dc38e8acc1WORKDIR
/opt/rosetta
05/08/2024 4:59 AM UTC
sha256:4da0f1f0b1c97a480d4efd295da0214ddd5ba9de61195cf9f654ecae50c15c81RUN
RUN |5 NVIDIA_BUILD_ID=92034664 GIT_USER_EMAIL=jax@nvidia.com GIT_USER_NAME=NVIDIA UPDATE_PATCHES=false UPDATED_TE_REF= /bin/sh -c <<"EOF" bash -exu MANIFEST_DIR=$(dirname ${MANIFEST_FILE}) if [[ "${UPDATE_PATCHES}" != "true" && "${UPDATE_PATCHES}" != "false" ]]; then echo "UPDATE_PATCHES can only be true or false" exit 1 fi if [[ "${UPDATE_PATCHES}" == "true" ]]; then cp -r /mnt/jax-toolbox/.github/container/patches ${MANIFEST_DIR}/ cp /mnt/jax-toolbox/.github/container/manifest.yaml ${MANIFEST_DIR}/manifest.yaml cp /mnt/jax-toolbox/.github/container/create-distribution.sh ${MANIFEST_DIR}/create-distribution.sh # TODO: remove cp /mnt/jax-toolbox/.github/container/pip-finalize.sh /usr/local/bin/ fi cp -r /mnt/jax-toolbox/rosetta /opt/rosetta if [[ -n "${UPDATED_TE_REF}" ]]; then TE_INSTALL_DIR=/opt/transformer-engine yq e ".transformer-engine.latest_verified_commit = \"${UPDATED_TE_REF}\"" -i $MANIFEST_FILE # Install from source instead of pre-built wheel sed -i -E 's@( file:///opt/transformer-engine)/dist/[^ ]*@\1@' /opt/pip-tools.d/requirements-te.in git -C $TE_INSTALL_DIR fetch -a if [[ "${UPDATED_TE_REF}" =~ ^pull/ ]]; then PR_ID=$(cut -d/ -f2 <<<"${UPDATED_TE_REF}") git -C $TE_INSTALL_DIR fetch origin ${UPDATED_TE_REF}:PR-${PR_ID} git -C $TE_INSTALL_DIR checkout PR-${PR_ID} else git -C $TE_INSTALL_DIR checkout ${UPDATED_TE_REF} fi fi # Setting the username/email is required to author commits from patches git config --global user.email "${GIT_USER_EMAIL}" git config --global user.name "${GIT_USER_NAME}" bash ${MANIFEST_DIR}/create-distribution.sh \
--manifest ${MANIFEST_FILE} \
--package paxml bash ${MANIFEST_DIR}/create-distribution.sh \
--manifest ${MANIFEST_FILE} \
--package praxis # Remove .gitconfig to avoid end-user authoring commits as the "build user" rm -f ~/.gitconfig EOF # buildkit
05/08/2024 4:59 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
NVTE_FUSED_ATTN=1
05/08/2024 4:59 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
ENABLE_TE=1
05/08/2024 4:59 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
UPDATED_TE_REF
05/08/2024 4:59 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
UPDATE_PATCHES
05/08/2024 4:59 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
GIT_USER_NAME
05/08/2024 4:59 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
GIT_USER_EMAIL
05/08/2024 4:59 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
NVIDIA_BUILD_ID=92034664
05/08/2024 4:59 AM UTC
...