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
06a6019a2797654a95a8ca349798e9231a7adf9b355bb680057f2eb208dd6940CONFIG
Entrypoint /opt/nvidia/nvidia_entrypoint.sh
10/20/2023 8:42 PM UTC
d682f49282a01500be6482a213c6d287e7981904872c78ae79a500d28e8204d0COPY
/opt/jax-source/.git /opt/jax-source/.git
10/20/2023 8:42 PM UTC
35e8fc99034206c723d9151d2646c27d79200f6bc3e5d87a3470179c201064e2ADD
build-jax.sh local_cuda_arch test-jax.sh /usr/local/bin/
10/20/2023 8:42 PM UTC
3b42e6197a7b352f7d85e07ad0cf4c2f737917b70e91c5d5b546caf52c2ba0dbARG
SRC_PATH_XLA
10/20/2023 8:42 PM UTC
de3e449cfd663d4d06e1cf8da68d28a99c9bfb17302f3166735d265cdee4a005ARG
SRC_PATH_JAX
10/20/2023 8:42 PM UTC
e54a9b13f57c211f8e666eb348dd08c81738104c55bcea82ab691ba06e0c17deRUN
SRC_PATH_JAX=/opt/jax-source SRC_PATH_XLA=/opt/xla-source BUILD_DATE=2023-10-20 { pip install flax ||
  true; } &&
  rm -rf ~/.cache/*
10/20/2023 8:42 PM UTC
bddf2aa1deaf5408bfaebe50f8c346f81d18daed447af4912a545751ceaf010aRUN
SRC_PATH_JAX=/opt/jax-source SRC_PATH_XLA=/opt/xla-source BUILD_DATE=2023-10-20 pip --disable-pip-version-check install ${SRC_PATH_JAX}/dist/*.whl &&
  pip --disable-pip-version-check install -e ${SRC_PATH_JAX} &&
  rm -rf ~/.cache/pip/
10/20/2023 8:42 PM UTC
9bdb7fdfee651ca484f38b686370ce1595b54debf2edc96ac5130f083317a920COPY
/opt/xla-source-no-git /opt/xla-source
10/20/2023 8:42 PM UTC
851b701d940f683f1cca7d1c40873fc95f08e8843ab77aaa3cba6382db1c2ffdCOPY
/opt/jax-source-no-git /opt/jax-source
10/20/2023 8:42 PM UTC
6fddf37a3d96105052a969aff9792928d26b23d6b9f722c297db13b7804ec7d0ENV
NCCL_NVLS_ENABLE=0
10/20/2023 8:42 PM UTC
...