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
f8cecb50510e80b421ac944ee4ec8b7f1bb85f444fa3e80e0956fa8bce5034d0CONFIG
Entrypoint /opt/nvidia/nvidia_entrypoint.sh
10/23/2024 6:31 AM UTC
877af25992fec0e2ce07be5d6740d5f4e1dd3e7cb638868af5c67f08c5f3575bRUN
RUN |9 URLREF_FLAX=https://github.com/google/flax.git#0e5a0b4de10c352c8374d2b611ad77f64ec9afc6 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE= BUILD_PATH_JAXLIB=/opt/jaxlibs NVIDIA_PRODUCT_NAME=JAX NVIDIA_JAX_VERSION=24.10 /bin/sh -c <<"EOF" bash -ex ls ${SRC_PATH_TRANSFORMER_ENGINE}/dist/*.whl echo "transformer-engine @ file://$(ls ${SRC_PATH_TRANSFORMER_ENGINE}/dist/*.whl)" > /opt/pip-tools.d/requirements-te.in EOF # buildkit
10/23/2024 6:28 AM UTC
ee6da40199124eedf81f5cb5f3b3600ab40fe50e29f28bfd9e0bb5f3c98f2554COPY
/opt/transformer-engine /opt/transformer-engine
10/23/2024 6:28 AM UTC
94735ddcc852f6a3d132ca9a6de05545ee48091203bce592c22e3610649ef718ENV
SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine
10/23/2024 6:28 AM UTC
51700c4adad5d9fe3e7519099e4aca484f111e5ae32b6ff092ac421024881508ENV
NVTE_FRAMEWORK=jax
10/23/2024 6:28 AM UTC
5d897a953f3e3b635eac2d321af89103e8538c41e02fcf51290cd71d6c9d11d9RUN
RUN |9 URLREF_FLAX=https://github.com/google/flax.git#0e5a0b4de10c352c8374d2b611ad77f64ec9afc6 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE= BUILD_PATH_JAXLIB=/opt/jaxlibs NVIDIA_PRODUCT_NAME=JAX NVIDIA_JAX_VERSION=24.10 /bin/sh -c <<"EOF" bash -ex git-clone.sh ${URLREF_FLAX} ${SRC_PATH_FLAX} sed -i "s|optax|optax>=0.2.3|" ${SRC_PATH_FLAX}/pyproject.toml sed -i "s|orbax-checkpoint|orbax-checkpoint<0.7.0|" ${SRC_PATH_FLAX}/pyproject.toml echo "-e file://${SRC_PATH_FLAX}" >> /opt/pip-tools.d/requirements-flax.in EOF # buildkit
10/23/2024 6:28 AM UTC
e1bf32a9bf086a3ae7b20701c2098bb4c55d00b39715cdb61c58fc44d5f4bd69RUN
RUN |9 URLREF_FLAX=https://github.com/google/flax.git#0e5a0b4de10c352c8374d2b611ad77f64ec9afc6 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE= BUILD_PATH_JAXLIB=/opt/jaxlibs NVIDIA_PRODUCT_NAME=JAX NVIDIA_JAX_VERSION=24.10 /bin/sh -c <<"EOF" bash -ex for component in $(ls ${BUILD_PATH_JAXLIB}); do echo "-e file://${BUILD_PATH_JAXLIB}/${component}" >> /opt/pip-tools.d/requirements-jax.in; done echo "-e file://${SRC_PATH_JAX}" >> /opt/pip-tools.d/requirements-jax.in echo "numpy<2.0.0" >> /opt/pip-tools.d/requirements-jax.in EOF # buildkit
10/23/2024 6:28 AM UTC
0d2d58d9111cbafaa8c2d9a307aafee64fd4d86f918e5988ac2a31dc25b7a487RUN
URLREF_FLAX=https://github.com/google/flax.git#0e5a0b4de10c352c8374d2b611ad77f64ec9afc6 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE= BUILD_PATH_JAXLIB=/opt/jaxlibs NVIDIA_PRODUCT_NAME=JAX NVIDIA_JAX_VERSION=24.10 mkdir -p /opt/pip-tools.d
10/23/2024 6:28 AM UTC
b8bf0ea8209222026914360356a3d662eac70432e1a02d78a98b243e99a8ced4ADD
build-jax.sh local_cuda_arch test-jax.sh /usr/local/bin/
10/23/2024 6:28 AM UTC
f4c6821e403bc6452062b1727502f47c630fa62f4aee13225cef6d3f7bc77fa3COPY
/opt/manifest.d/git-clone.yaml /opt/manifest.d/git-clone.yaml
10/23/2024 6:28 AM UTC
...

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.