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:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
NVIDIA_JAX_VERSION=26.04
04/23/2026 5:10 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
NVIDIA_PRODUCT_NAME=JAX
04/23/2026 5:10 PM UTC
sha256:c6e4b36e806d1847cff318015c1e1377016b382aed4b907d569bd333858c78feRUN
URLREF_FLAX=https://github.com/google/flax.git#1fcc2a5c7a076a47a1dfe7b82825f8c007b59ed7 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_NSYS_JAX=/opt/nsys_jax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2026-04-23 BUILD_PATH_JAXLIB=/opt/jaxlibs pip-finalize.sh
04/23/2026 5:10 PM UTC
sha256:d10f98e3ef3838fba106ff77cd5977c0201cf45d84c8820fe4929e79fdc057eeCOPY
pinned/26.04-devel/jax-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
04/23/2026 5:09 PM UTC
sha256:2496fcc20b6e04e19febe83d5cda0af20bdf9da25e6092c14eceaeac3243bdecRUN
URLREF_FLAX=https://github.com/google/flax.git#1fcc2a5c7a076a47a1dfe7b82825f8c007b59ed7 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_NSYS_JAX=/opt/nsys_jax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2026-04-23 BUILD_PATH_JAXLIB=/opt/jaxlibs echo "cuda-tile==${CUTILE_PYTHON_VERSION}" > /opt/pip-tools.d/requirements-cuda-tile.in
04/23/2026 5:08 PM UTC
sha256:1061fc49a5178e4329bc472cc5b2a76af94685dde77563fd25255c5cac59410fRUN
URLREF_FLAX=https://github.com/google/flax.git#1fcc2a5c7a076a47a1dfe7b82825f8c007b59ed7 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_NSYS_JAX=/opt/nsys_jax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2026-04-23 BUILD_PATH_JAXLIB=/opt/jaxlibs install-nsys-jax.sh ${SRC_PATH_NSYS_JAX}
04/23/2026 5:08 PM UTC
sha256:3514e9c393440bd66082303e321cf586beb9aae5b358aa1442d614dc1d2fda09ADD
install-nsys-jax.sh /usr/local/bin
04/23/2026 5:08 PM UTC
sha256:13336517bd202f06e7fe50a42f9407206d5897ddf65e9ff26b000faae8bb4f74COPY
nsys_jax/ /opt/nsys_jax
04/23/2026 5:08 PM UTC
sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
SRC_PATH_NSYS_JAX=/opt/nsys_jax
04/23/2026 5:08 PM UTC
sha256:b1efd299f774981adb32a6f5cf12402c194bcc151c33471838f6f53ea0b6708eRUN
RUN |8 URLREF_FLAX=https://github.com/google/flax.git#1fcc2a5c7a076a47a1dfe7b82825f8c007b59ed7 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_NSYS_JAX=/opt/nsys_jax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2026-04-23 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/23/2026 5:08 PM UTC
...