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
ca6f23cb64ed434d3e67090cfb6eafd7c81215079e3db814134ae8a11cb021deCONFIG
Entrypoint /opt/nvidia/nvidia_entrypoint.sh
08/25/2023 10:01 PM UTC
6810de8b8b1c2e8f188468eb1d3863b0ef50494e3c16eee01d0baafa1b522796COPY
/opt/jax-source/.git /opt/jax-source/.git
08/25/2023 10:01 PM UTC
79fe3c6abe8cdac884e9bf06060fe3f47a4d0b549362e2f2568aa377d93652daADD
build-jax.sh local_cuda_arch test-jax.sh /usr/local/bin/
08/25/2023 10:01 PM UTC
bc2b4857c24fc861398726fb18bbd3df5864719c25ccaa78a6acf3d02207ade1ARG
SRC_PATH_XLA
08/25/2023 10:01 PM UTC
44121880f12d35c73ed7ccb92b1ed86d749627af1e0fca108d2d5eaff6ef4a9bARG
SRC_PATH_JAX
08/25/2023 10:01 PM UTC
038d6e27529fe99f36ec0a6df84421429dbe0bde5dc0a50445322723a045c6e7RUN
SRC_PATH_JAX=/opt/jax-source SRC_PATH_XLA=/opt/xla-source BUILD_DATE=2023-08-25 { pip install flax ||
  true; } &&
  rm -rf ~/.cache/pip
08/25/2023 10:01 PM UTC
4cb9b6d65d58edd249e496f682efed55a2a57fba6860d974dd9b1c90ecf671fcRUN
SRC_PATH_JAX=/opt/jax-source SRC_PATH_XLA=/opt/xla-source BUILD_DATE=2023-08-25 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/
08/25/2023 10:01 PM UTC
abf375cb7415cfb8df2c23961616fd91bfffe269ff933acf716cbd8d1fef02a0COPY
/opt/xla-source-no-git /opt/xla-source
08/25/2023 10:01 PM UTC
b64282f303b5caf52b225639fa32480d2dd219797d5f7efc8b15c47552b8b884COPY
/opt/jax-source-no-git /opt/jax-source
08/25/2023 10:01 PM UTC
7ed78ad4bc6415c4188a0944c804b9d8f05e0dc027c1eb0f47d3414eaf3840feENV
NCCL_IB_SL=1
08/25/2023 10:01 PM 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.