Skip to main content
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
JAX
Container
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:1fc014731cac66b26c51e386b4dea9e30f0d8920ce9298fc6325165ac2e92b9cRUN
RUN /bin/sh -c <<"EOF" bash -e -x git -C /opt/jax-toolbox-mirror show origin/t5x-23.10-release:rosetta/rosetta/projects/t5x/README.md > /opt/rosetta/rosetta/projects/t5x/README.md EOF # buildkit
10/25/2023 4:49 AM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
CUDA_MODULE_LOADING=EAGER
10/25/2023 4:49 AM UTC
sha256:98ad8d13d1a0732d5af5c04cde91557de592402531b75c4afda60c14b38214f7ADD
test-vit.sh /usr/local/bin
10/20/2023 7:09 PM UTC
sha256:2b25a64597b6a7dbab787cd92f4512da1a07f212749ce0844c9967d5f2971b6dRUN
RUN |2 GIT_USER_EMAIL=jax@nvidia.com GIT_USER_NAME=NVIDIA /bin/sh -c <<EOF bash -e pip install -e /opt/rosetta --extra-index-url https://developer.download.nvidia.com/compute/redist rm -rf ~/.cache/pip/ EOF # buildkit
10/20/2023 7:09 PM UTC
sha256:4f4fb700ef54461cfa02571ae0db9a0dc1e0cdb5577484a6d75e68dc38e8acc1WORKDIR
/opt/rosetta
10/20/2023 7:09 PM UTC
sha256:c029cc89262e00152d799c36568aba7488d1124210d16a9a54d631d09768081dRUN
RUN |2 GIT_USER_EMAIL=jax@nvidia.com GIT_USER_NAME=NVIDIA /bin/sh -c <<"EOF" bash -e bash create-distribution.sh \
-p patchlist-t5x.txt \
-m https://github.com/nvjax-svc-0/t5x.git \
-d $(dirname $(python -c "import t5x; print(*t5x.__path__)")) \
-e /opt/t5x-mirror bash create-distribution.sh \
-p patchlist-flax.txt \
-m https://github.com/nvjax-svc-0/flax.git \
-d $(dirname $(python -c "import flax; print(*flax.__path__)")) \
-e /opt/flax-mirror rm -rf $(find /opt -name "__pycache__") EOF # buildkit
10/20/2023 7:09 PM UTC
sha256:4f4fb700ef54461cfa02571ae0db9a0dc1e0cdb5577484a6d75e68dc38e8acc1WORKDIR
/opt/rosetta
10/20/2023 7:09 PM UTC
sha256:8c18990ec90412922449215da0947171188a02d17cdac8a57f2d839e89bc7a95COPY
/ /opt/rosetta
10/20/2023 7:09 PM UTC
sha256:205aa327cd304b32b420a3c153efda959cc08c7eed15a8d7d55bb3e787747f16RUN
RUN |2 GIT_USER_EMAIL=jax@nvidia.com GIT_USER_NAME=NVIDIA /bin/sh -c <<EOF bash -e git config --global user.email "${GIT_USER_EMAIL}" git config --global user.name "${GIT_USER_NAME}" EOF # buildkit
10/20/2023 7:09 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
GIT_USER_NAME
10/20/2023 7:09 PM UTC
1-10 of 106 items

Modal Content

NVIDIA uses cookies to improve your experience on our web site. We and our third-party partners also use cookies and other tools to collect and record information you provide as well as information about your interactions with our websites for performance improvement, analytics, and to assist in marketing efforts. By clicking "Accept All", you consent to our use of cookies and other tools as described in our Cookie Policy. You can manage your cookie settings by clicking on "Manage Settings." By continuing to use this site or by clicking one of the buttons below, you agree to our Terms of Service (which contains important waivers). Please see our Privacy Policy for more information on our privacy practices.

Manage Settings