mirror of
https://github.com/ROCm/jax.git
synced 2025-04-24 17:56:23 +00:00
94f97b920f
Most of the work here is porting the LAPACK interface from Cython to plain C++. This is something I wanted to do anyway to make use of C++ templating facilities: the code is noticeably shorter in C++. This change removes the only use of Cython in JAX. It also removes the need for a build-time dependency on Scipy, which we only needed for Cython cimport reasons. When using C++, we most likely do not want to fetch LAPACK and BLAS kernels from Python. Therefore we add another option: we define the LAPACK functions we need using weak symbols where supported; the user can then simply link against LAPACK to provide the necessary symbols. Added a jaxlib:cpu_kernels module to facilitate using the JAX CPU kernels from C++. PiperOrigin-RevId: 394705605
36 lines
1.5 KiB
Docker
36 lines
1.5 KiB
Docker
FROM gcr.io/tensorflow-testing/nosla-cuda10.0-cudnn7-ubuntu16.04-manylinux2010
|
|
LABEL maintainer "Matt Johnson <mattjj@google.com>"
|
|
|
|
WORKDIR /
|
|
# TODO(skyewm): delete the following line when no longer necessary.
|
|
RUN rm -f /etc/apt/sources.list.d/jonathonf-ubuntu-python-3_6-xenial.list
|
|
RUN apt-get update
|
|
RUN apt-get install libffi-dev
|
|
RUN git clone --branch v1.2.21 https://github.com/pyenv/pyenv.git /pyenv
|
|
ENV PYENV_ROOT /pyenv
|
|
RUN /pyenv/bin/pyenv install 3.7.2
|
|
RUN /pyenv/bin/pyenv install 3.8.0
|
|
RUN /pyenv/bin/pyenv install 3.9.0
|
|
|
|
# We pin numpy to the minimum permitted version to avoid compatibility issues.
|
|
RUN eval "$(/pyenv/bin/pyenv init -)" && /pyenv/bin/pyenv local 3.7.2 && pip install numpy==1.18.5 setuptools wheel six auditwheel
|
|
RUN eval "$(/pyenv/bin/pyenv init -)" && /pyenv/bin/pyenv local 3.8.0 && pip install numpy==1.18.5 setuptools wheel six auditwheel
|
|
RUN eval "$(/pyenv/bin/pyenv init -)" && /pyenv/bin/pyenv local 3.9.0 && pip install numpy==1.19.4 setuptools wheel six auditwheel
|
|
|
|
# Change the CUDA version if it doesn't match the installed version.
|
|
ARG JAX_CUDA_VERSION=10.0
|
|
COPY install_cuda.sh /install_cuda.sh
|
|
RUN chmod +x /install_cuda.sh
|
|
RUN /bin/bash -c 'if [[ ! "$CUDA_VERSION" =~ ^$JAX_CUDA_VERSION.*$ ]]; then \
|
|
/install_cuda.sh $JAX_CUDA_VERSION; \
|
|
fi'
|
|
|
|
|
|
WORKDIR /
|
|
COPY build_wheel_docker_entrypoint.sh /build_wheel_docker_entrypoint.sh
|
|
RUN chmod +x /build_wheel_docker_entrypoint.sh
|
|
|
|
WORKDIR /build
|
|
ENV TEST_TMPDIR /build
|
|
ENTRYPOINT ["/build_wheel_docker_entrypoint.sh"]
|