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:12c11b15227af8ac23ff4088c376df4e81eb0c8bc78dda78d51618662db630dcRUN
    RUN /bin/sh -c <<"EOF" bash -e -x pushd /opt/paxml git -C /opt/pax-mirror diff $(git -C /opt/pax-mirror merge-base origin/main origin/patch/23.10_release)..origin/patch/23.10_release | git apply git add -u git commit -m "23.10 Patches" popd git -C /opt/jax-toolbox-mirror show origin/pax-23.10-release:rosetta/rosetta/projects/pax/README.md > /opt/rosetta/rosetta/projects/pax/README.md git -C /opt/jax-toolbox-mirror show origin/pax-23.10-release:rosetta/rosetta/projects/pax/scripts/example_slurm_pile.sub > /opt/rosetta/rosetta/projects/pax/scripts/example_slurm_pile.sub EOF # buildkit
    10/25/2023 6:02 PM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
    ENABLE_TE=1
    10/25/2023 6:02 PM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ENV
    CUDA_MODULE_LOADING=EAGER
    10/25/2023 6:02 PM UTC
    sha256:92fbfbaee18a5ae80f615252ca6e81e494a2c171918678ac60725612d1bc6e93RUN
    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
    sha256:4f4fb700ef54461cfa02571ae0db9a0dc1e0cdb5577484a6d75e68dc38e8acc1WORKDIR
    /opt/rosetta
    10/12/2023 4:55 PM UTC
    sha256:ee38ba970e22319ffbea3a0da747d44924d98b466cd643e9aa1fe58614815585COPY
    / /opt/rosetta
    10/12/2023 4:55 PM UTC
    sha256:2a4c1db0aa87620daebdc3eeb8629584508e5f04404bc7f30df209073d00b0abRUN
    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
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
    GIT_USER_NAME
    10/12/2023 4:55 PM UTC
    sha256:a3ed95caeb02ffe68cdd9fd84406680ae93d633cb16422d00e8a7c22955b46d4ARG
    GIT_USER_EMAIL
    10/12/2023 4:55 PM UTC
    sha256:afcf3819105c0cb9ef1511f36d59e57e891878748bd7a5a4cd5208274f66f4c5ADD
    test-pax.sh /usr/local/bin
    10/12/2023 4:34 PM UTC
    ...