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
1663075a0ac733e8967bc2809390e279fe4a56d1c0126d6c38937d6285aa544fCONFIG
Entrypoint /opt/nvidia/nvidia_entrypoint.sh
04/16/2025 9:44 PM UTC
ee191f680274b0a4932369647c7276515758ac27d4a6944ace0e1330e4495a2cENV
NVIDIA_PRODUCT_NAME=JAX
04/16/2025 9:44 PM UTC
90749c8bcfa569d400791c70307123bb59345cd486d2e3f91f8591cc9f21893cRUN
URLREF_FLAX=https://github.com/google/flax.git#v0.10.5 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2025-04-16 BUILD_PATH_JAXLIB=/opt/jaxlibs pip-finalize.sh
04/16/2025 9:44 PM UTC
6eded15089ce9a219505a68a72fa0d77bbcf4a4db81256a32b875b3a88da795bRUN
RUN |7 URLREF_FLAX=https://github.com/google/flax.git#v0.10.5 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2025-04-16 BUILD_PATH_JAXLIB=/opt/jaxlibs /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
04/16/2025 9:40 PM UTC
0e7b5de94925123ce72b1498986437e5131cd33b558cc203cb5777977621abf0COPY
/opt/transformer-engine /opt/transformer-engine
04/16/2025 9:40 PM UTC
660585f9f887188334a256b46643ccd880bd11cfe076b591fc8c82e80828f0f2ENV
SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine
04/16/2025 9:40 PM UTC
bf92b7074f8e7dd9c1e068c9f445bc0638859270ac2aadc3aa10bcfd2116dbfeENV
NVTE_FRAMEWORK=jax
04/16/2025 9:40 PM UTC
4301bcc956ef2e7edd8b2848eaedf692ac2974b5ad7fb3477790dddeca29ab94RUN
RUN |7 URLREF_FLAX=https://github.com/google/flax.git#v0.10.5 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2025-04-16 BUILD_PATH_JAXLIB=/opt/jaxlibs /bin/sh -c <<"EOF" bash -ex git-clone.sh ${URLREF_FLAX} ${SRC_PATH_FLAX} echo "-e file://${SRC_PATH_FLAX}" >> /opt/pip-tools.d/requirements-flax.in EOF # buildkit
04/16/2025 9:40 PM UTC
4418e68b78632f050c8210e2ac6513ebcfc0bd0d3fae0d7289d7bebcdb8ea9abRUN
RUN |7 URLREF_FLAX=https://github.com/google/flax.git#v0.10.5 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2025-04-16 BUILD_PATH_JAXLIB=/opt/jaxlibs /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}[k8s]" >> /opt/pip-tools.d/requirements-jax.in EOF # buildkit
04/16/2025 9:40 PM UTC
2fb095ec0c3b56ad9e2fcafdf52d9849d914fa4a4c41808d77264affceca5164RUN
URLREF_FLAX=https://github.com/google/flax.git#v0.10.5 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2025-04-16 BUILD_PATH_JAXLIB=/opt/jaxlibs mkdir -p /opt/pip-tools.d
04/16/2025 9:40 PM UTC
...