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

Update jaxlib build scripts to build Python 3.8.0 wheels. (#1612)

Browse Source
This commit is contained in:
Peter Hawkins 2019-10-31 11:46:37 -04:00 committed by GitHub
parent 33e5d07304
commit b3e4a1a850
No known key found for this signature in database
GPG Key ID: 4AEE18F83AFDEB23
3 changed files with 24 additions and 13 deletions
Show Stats Download Patch File Download Diff File Expand all files Collapse all files

16
build/Dockerfile
View File

@ -2,18 +2,24 @@ FROM gcr.io/tensorflow-testing/nosla-cuda10.0-cudnn7-ubuntu16.04-manylinux2010
LABEL maintainer "Matt Johnson <mattjj@google.com>"
WORKDIR /
RUN git clone https://github.com/pyenv/pyenv.git /pyenv
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
RUN /pyenv/bin/pyenv install 3.5.6
RUN /pyenv/bin/pyenv install 3.6.8
RUN /pyenv/bin/pyenv install 3.7.2
RUN /pyenv/bin/pyenv install 3.8.0
# Install build-dependencies of scipy.
# TODO(phawkins): remove when there are scipy wheels for Python 3.8.
RUN apt-get update && apt-get install -y libopenblas-dev gfortran
# We pin numpy to a version < 1.16 to avoid version compatibility issues.
RUN eval "$(/pyenv/bin/pyenv init -)" && /pyenv/bin/pyenv local 2.7.15 && pip install numpy==1.15.4 scipy cython setuptools wheel future
RUN eval "$(/pyenv/bin/pyenv init -)" && /pyenv/bin/pyenv local 3.5.6 && pip install numpy==1.15.4 scipy cython setuptools wheel
RUN eval "$(/pyenv/bin/pyenv init -)" && /pyenv/bin/pyenv local 3.6.8 && pip install numpy==1.15.4 scipy cython setuptools wheel
RUN eval "$(/pyenv/bin/pyenv init -)" && /pyenv/bin/pyenv local 3.7.2 && pip install numpy==1.15.4 scipy cython setuptools wheel
RUN eval "$(/pyenv/bin/pyenv init -)" && /pyenv/bin/pyenv local 2.7.15 && pip install numpy==1.15.4 scipy cython setuptools wheel future six
RUN eval "$(/pyenv/bin/pyenv init -)" && /pyenv/bin/pyenv local 3.5.6 && pip install numpy==1.15.4 scipy cython setuptools wheel six
RUN eval "$(/pyenv/bin/pyenv init -)" && /pyenv/bin/pyenv local 3.6.8 && pip install numpy==1.15.4 scipy cython setuptools wheel six
RUN eval "$(/pyenv/bin/pyenv init -)" && /pyenv/bin/pyenv local 3.7.2 && pip install numpy==1.15.4 scipy cython setuptools wheel six
RUN eval "$(/pyenv/bin/pyenv init -)" && /pyenv/bin/pyenv local 3.8.0 && pip install numpy==1.17.3 scipy cython setuptools wheel six
# Change the CUDA version if it doesn't match the installed version.
ARG JAX_CUDA_VERSION=10.0

2
build/build_jaxlib_wheels.sh
View File

@ -1,7 +1,7 @@
#!/bin/bash
set -xev
PYTHON_VERSIONS="2.7.15 3.5.6 3.6.8 3.7.2"
PYTHON_VERSIONS="2.7.15 3.5.6 3.6.8 3.7.2 3.8.0"
CUDA_VERSIONS="9.0 9.2 10.0 10.1"
CUDA_VARIANTS="cuda" # "cuda-included"

19
build/build_jaxlib_wheels_macos.sh
View File

@ -19,15 +19,19 @@ PLATFORM_TAG="macosx_10_9_x86_64"
build_jax () {
PY_VERSION="$1"
PY_TAG="$2"
NUMPY_VERSION="$3"
SCIPY_VERSION="$4"
echo "\nBuilding JAX for Python ${PY_VERSION}, tag ${PY_TAG}"
echo "NumPy version ${NUMPY_VERSION}, SciPy version ${SCIPY_VERSION}"
pyenv install -s "${PY_VERSION}"
VENV="jax-build-${PY_VERSION}"
pyenv virtualenv-delete -f "${VENV}"
pyenv virtualenv "${PY_VERSION}" "${VENV}"
pyenv activate "${VENV}"
# We pin the Numpy wheel to a version < 1.16.0, because Numpy extensions built
# at 1.16.0 are not backward compatible to earlier Numpy versions.
pip install numpy==1.15.4 scipy==1.2.0 wheel future
# We pin the Numpy wheel to a version < 1.16.0 for Python releases prior to
# 3.8, because Numpy extensions built at 1.16.0 are not backward compatible to
# earlier Numpy versions.
pip install numpy==$NUMPY_VERSION scipy==$SCIPY_VERSION wheel future six
rm -fr build/build
python build/build.py
cd build
@ -39,7 +43,8 @@ build_jax () {
rm -fr build/dist
build_jax 2.7.15 cp27
build_jax 3.5.6 cp35
build_jax 3.6.8 cp36
build_jax 3.7.2 cp37
build_jax 2.7.15 cp27 1.15.4 1.2.0
build_jax 3.5.6 cp35 1.15.4 1.2.0
build_jax 3.6.8 cp36 1.15.4 1.2.0
build_jax 3.7.2 cp37 1.15.4 1.2.0
build_jax 3.8.0 cp38 1.17.3 1.3.1