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
14066da5e7f171f6fac7dbe624d6f82af4c75610919342c0f7b9d1e978102417CONFIG
Entrypoint /opt/nvidia/nvidia_entrypoint.sh; WorkingDir /opt/rosetta
10/25/2023 6:02 PM UTC
dee4fa9e0e31d112e772f28a0508bf135ea13a3e283f31b21c87b5eb6dbb1ef7ENV
ENABLE_TE=1
10/25/2023 6:02 PM UTC
5329bb2ab32b440cb74615a660473543819e1a9a26f59a8533feab71a86381f0ENV
CUDA_MODULE_LOADING=EAGER
10/25/2023 6:02 PM UTC
de0f6ab9aff33673ff05bd15c7dd1a969453cbdfd401f392198272c54d485972RUN
RUN |2 GIT_USER_EMAIL=jax@nvidia.com GIT_USER_NAME=NVIDIA /bin/sh -c <<"EOF" bash -e bash create-distribution.sh \
-p patchlist-paxml.txt \
-m https://github.com/nvjax-svc-0/paxml.git \
-d $(dirname $(python -c "import paxml; print(*paxml.__path__)")) \
-e /opt/pax-mirror bash create-distribution.sh \
-p patchlist-praxis.txt \
-m https://github.com/nvjax-svc-0/praxis.git \
-d $(dirname $(python -c "import praxis; print(*praxis.__path__)")) \
-e /opt/praxis-mirror rm -rf $(find /opt -name "__pycache__") EOF # buildkit
10/12/2023 4:55 PM UTC
fcce817fd56544598c4404baa915a7d1e2f442a417c9d3acb5623ba85be4ac12WORKDIR
/opt/rosetta
10/12/2023 4:55 PM UTC
72e514da0d2820d7742e1cc8bd9a1fc220d64f81edeb371feca369a45509df2fCOPY
/ /opt/rosetta
10/12/2023 4:55 PM UTC
3d69ea93782e20568ab5a95f1f1e071e69769da3d49deb117e20b50669dd43beRUN
RUN |2 GIT_USER_EMAIL=jax@nvidia.com GIT_USER_NAME=NVIDIA /bin/sh -c <<EOF bash -e git config --global user.email "${GIT_USER_EMAIL}" git config --global user.name "${GIT_USER_NAME}" EOF # buildkit
10/12/2023 4:55 PM UTC
4ab8ed275834fc443e52a7cd69bcb283bbac6caeb9c42fb63186e9d7f4c0fe6bARG
GIT_USER_NAME
10/12/2023 4:55 PM UTC
d2b3e3ac6a951b41796156d401cb6bbc8f694172697a0cca82b654ecba033b46ARG
GIT_USER_EMAIL
10/12/2023 4:55 PM UTC
36c016255ac85f3dbf208914e0736705d447aee1467878c4a2d34cf03ffd0c3dADD
test-pax.sh /usr/local/bin
10/12/2023 4:34 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.