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
6d6756d60f1cb742dbd784389f23bd754b7ed4c8199f2016d9b0b0ef0b7af3abCONFIG
Entrypoint /opt/nvidia/nvidia_entrypoint.sh
02/26/2026 12:30 AM UTC
b1a5d436043d535c112284077a9df8ebdea04ecbbe35113c0fa84505699eeb40ENV
NVIDIA_PRODUCT_NAME=JAX
02/26/2026 12:30 AM UTC
3d2855e02dc46569d382afe4818fcb089cea260666883b0824539951ea5ffb39RUN
URLREF_FLAX=https://github.com/google/flax.git#7e5c70a0d119c2b4f3c9465189e9c10f00263471 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2026-02-25 BUILD_PATH_JAXLIB=/opt/jaxlibs pip-finalize.sh
02/26/2026 12:30 AM UTC
5e5708180493c858251482d3f8cf4b837ad5d7cf28e0c0880f873a5193eb1f34COPY
pinned/26.02-devel/jax-amd64.txt /opt/pip-tools.d/requirements-pinned.txt
02/26/2026 12:30 AM UTC
145626d9ba50e11f34073c69bfb2fad34c2b9aded7a5572f98121ab5889f197fRUN
RUN |7 URLREF_FLAX=https://github.com/google/flax.git#7e5c70a0d119c2b4f3c9465189e9c10f00263471 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2026-02-25 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
02/26/2026 12:28 AM UTC
8127baca6e510237599becfbb30d42fe89da15de690d541cae8e051064c7122aCOPY
/opt/transformer-engine /opt/transformer-engine
02/26/2026 12:28 AM UTC
bf68956d557c3cdd7688eed80930e0c8833b5ccc09fe3860515deb0ee548b654ENV
SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine
02/26/2026 12:28 AM UTC
8f3b458bc5cc06283fdde26db9c26898866e36499a4c3ea8b982cc16a00d4c4dRUN
RUN |7 URLREF_FLAX=https://github.com/google/flax.git#7e5c70a0d119c2b4f3c9465189e9c10f00263471 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2026-02-25 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
02/26/2026 12:28 AM UTC
4bafc351f762b2bc13026bcb5c45ec8e18145da6cefba8f15c506070327ca4b0RUN
RUN |7 URLREF_FLAX=https://github.com/google/flax.git#7e5c70a0d119c2b4f3c9465189e9c10f00263471 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2026-02-25 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
02/26/2026 12:28 AM UTC
ff4ed668a7756ebf9842f2f8f62edd61087fd6840d4c507c622d4b83ab08b439RUN
URLREF_FLAX=https://github.com/google/flax.git#7e5c70a0d119c2b4f3c9465189e9c10f00263471 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_FLAX=/opt/flax SRC_PATH_TRANSFORMER_ENGINE=/opt/transformer-engine BUILD_DATE=2026-02-25 BUILD_PATH_JAXLIB=/opt/jaxlibs mkdir -p /opt/pip-tools.d
02/26/2026 12: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.