Entrypoint /opt/nvidia/nvidia_entrypoint.sh

RUN |4 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_TE=/opt/transformer-engine BUILD_DATE= /bin/sh -c <<"EOF" bash -ex -o pipefail pip install ninja &&
  rm -rf ~/.cache/pip get-source.sh -l transformer-engine -m ${MANIFEST_FILE} pushd ${SRC_PATH_TE} ## apply release-specific patches: cherry-pick changes after the last release tag git cherry-pick 8e672ff0758033c348e263dbcd6a4b3578c01161 git cherry-pick bfe21c3d68b0a9951e5716fb520045db53419c5e git cherry-pick 7c1828f80edc1405d4ef1a7780c9e0046beab5c7 # https://github.com/NVIDIA/TransformerEngine/pull/745 ## end of apply release-specific patches python setup.py bdist_wheel &&
  rm -rf build echo "transformer-engine @ file://$(ls ${SRC_PATH_TE}/dist/*.whl)" >> /opt/pip-tools.d/requirements-te.in EOF # buildkit

SRC_PATH_TE=/opt/transformer-engine

NVTE_FRAMEWORK=jax

SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_TE=/opt/transformer-engine BUILD_DATE= get-source.sh -l flax -m ${MANIFEST_FILE} -o /opt/pip-tools.d/requirements-flax.in

RUN |4 SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_TE=/opt/transformer-engine BUILD_DATE= /bin/sh -c <<"EOF" bash -ex # Encourage a newer numpy so that pip's dependency resolver will allow newer
# versions of other packages that rely on newer numpy, but also include fixes
# for compatibility with newer JAX versions. e.g. chex.
echo "numpy >= 1.24.1"                                  >> /opt/pip-tools.d/requirements-jax.in
echo "-e file://${SRC_PATH_JAX}"                        >> /opt/pip-tools.d/requirements-jax.in
echo "jaxlib @ file://$(ls ${SRC_PATH_JAX}/dist/*.whl)" >> /opt/pip-tools.d/requirements-jax.in
EOF # buildkit

SRC_PATH_JAX=/opt/jax SRC_PATH_XLA=/opt/xla SRC_PATH_TE=/opt/transformer-engine BUILD_DATE= mkdir -p /opt/pip-tools.d

build-jax.sh local_cuda_arch test-jax.sh /usr/local/bin/

/opt/xla /opt/xla

/opt/jax /opt/jax