NVIDIA
JAX
Container
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:1fc014731cac66b26c51e386b4dea9e30f0d8920ce9298fc6325165ac2e92b9cRUN
    RUN /bin/sh -c <<"EOF" bash -e -x git -C /opt/jax-toolbox-mirror show origin/t5x-23.10-release:rosetta/rosetta/projects/t5x/README.md > /opt/rosetta/rosetta/projects/t5x/README.md EOF # buildkit
    10/25/2023 4:49 AM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
    CUDA_MODULE_LOADING=EAGER
    10/25/2023 4:49 AM UTC
    sha256:98ad8d13d1a0732d5af5c04cde91557de592402531b75c4afda60c14b38214f7ADD
    test-vit.sh /usr/local/bin
    10/20/2023 7:09 PM UTC
    sha256:2b25a64597b6a7dbab787cd92f4512da1a07f212749ce0844c9967d5f2971b6dRUN
    RUN |2 GIT_USER_EMAIL=jax@nvidia.com GIT_USER_NAME=NVIDIA /bin/sh -c <<EOF bash -e pip install -e /opt/rosetta --extra-index-url https://developer.download.nvidia.com/compute/redist rm -rf ~/.cache/pip/ EOF # buildkit
    10/20/2023 7:09 PM UTC
    sha256:4f4fb700ef54461cfa02571ae0db9a0dc1e0cdb5577484a6d75e68dc38e8acc1WORKDIR
    /opt/rosetta
    10/20/2023 7:09 PM UTC
    sha256:c029cc89262e00152d799c36568aba7488d1124210d16a9a54d631d09768081dRUN
    RUN |2 GIT_USER_EMAIL=jax@nvidia.com GIT_USER_NAME=NVIDIA /bin/sh -c <<"EOF" bash -e bash create-distribution.sh \
    -p patchlist-t5x.txt \
    -m https://github.com/nvjax-svc-0/t5x.git \
    -d $(dirname $(python -c "import t5x; print(*t5x.__path__)")) \
    -e /opt/t5x-mirror bash create-distribution.sh \
    -p patchlist-flax.txt \
    -m https://github.com/nvjax-svc-0/flax.git \
    -d $(dirname $(python -c "import flax; print(*flax.__path__)")) \
    -e /opt/flax-mirror rm -rf $(find /opt -name "__pycache__") EOF # buildkit
    10/20/2023 7:09 PM UTC
    sha256:4f4fb700ef54461cfa02571ae0db9a0dc1e0cdb5577484a6d75e68dc38e8acc1WORKDIR
    /opt/rosetta
    10/20/2023 7:09 PM UTC
    sha256:8c18990ec90412922449215da0947171188a02d17cdac8a57f2d839e89bc7a95COPY
    / /opt/rosetta
    10/20/2023 7:09 PM UTC
    sha256:205aa327cd304b32b420a3c153efda959cc08c7eed15a8d7d55bb3e787747f16RUN
    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/20/2023 7:09 PM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
    GIT_USER_NAME
    10/20/2023 7:09 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.