mirrors/rocm_jax
1
0
Fork 0
mirror of https://github.com/ROCm/jax.git synced 2025-04-18 21:06:06 +00:00
Code Issues Packages Projects Releases Wiki Activity

jaxlib build improvements (#1742)

Browse Source
This commit is contained in:
Skye Wanderman-Milne 2019-11-21 16:25:24 -08:00 committed by GitHub
parent 27aa76e6a6
commit 39e17398ee
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
2 changed files with 4 additions and 2 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

2
build/Dockerfile
View File

@ -2,6 +2,8 @@ FROM gcr.io/tensorflow-testing/nosla-cuda10.0-cudnn7-ubuntu16.04-manylinux2010
LABEL maintainer "Matt Johnson <mattjj@google.com>"
WORKDIR /
RUN apt-get update
RUN apt-get install libffi-dev
RUN git clone --branch v1.2.14 https://github.com/pyenv/pyenv.git /pyenv
ENV PYENV_ROOT /pyenv
RUN /pyenv/bin/pyenv install 2.7.15

4
build/build_jaxlib_wheels.sh
View File

@ -13,7 +13,7 @@ for PYTHON_VERSION in $PYTHON_VERSIONS
do
mkdir -p dist/nocuda/
docker run -it --tmpfs /build:exec --rm -v $(pwd)/dist:/dist jaxbuild $PYTHON_VERSION nocuda
mv dist/*.whl dist/nocuda/
mv -f dist/*.whl dist/nocuda/
done
# build the cuda linux packages, tagging with linux_x86_64
@ -26,7 +26,7 @@ do
do
mkdir -p dist/${CUDA_VARIANT}${CUDA_VERSION//.}
docker run -it --tmpfs /build:exec --rm -v $(pwd)/dist:/dist jaxbuild $PYTHON_VERSION $CUDA_VARIANT
mv dist/*.whl dist/${CUDA_VARIANT}${CUDA_VERSION//.}/
mv -f dist/*.whl dist/${CUDA_VARIANT}${CUDA_VERSION//.}/
done
done
done