check out build files

build/Dockerfile

@@ -0,0 +1,21 @@
+ARG CUDA_VERSION=9.2
+FROM nvidia/cuda:$CUDA_VERSION-cudnn7-devel-ubuntu16.04
+LABEL maintainer "Matt Johnson <mattjj@google.com>"
+
+RUN apt-get update && apt-get install -y --no-install-recommends \
+            dh-autoreconf git curl \
+            python python-pip python-dev \
+            python3 python3-pip python3-dev
+RUN pip install numpy setuptools wheel && pip3 install numpy setuptools wheel
+
+RUN git clone https://github.com/nixos/patchelf /tmp/patchelf
+WORKDIR /tmp/patchelf
+RUN bash bootstrap.sh && ./configure && make && make install && rm -r /tmp/patchelf
+
+WORKDIR /
+RUN curl -O https://raw.githubusercontent.com/google/jax/744ac4821afafaff86b0f1820f05ac7c1186e276/build/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"]

build/build_wheel_docker_entrypoint.sh

@@ -0,0 +1,47 @@
+#!/bin/bash -xev
+if [ ! -d "/dist" ]
+then
+  echo "/dist must be mounted to produce output"
+  exit 1
+fi
+
+git clone -b binary-distros https://github.com/google/jax /build/jax
+cd /build/jax
+
+usage() {
+  echo "usage: ${0##*/} [python2|python3] [cuda-included|cuda|nocuda]"
+  exit 1
+}
+
+if [[ $# != 2 ]]
+then
+  usage
+fi
+
+case $1 in
+  py2)
+    update-alternatives --install /usr/bin/python python /usr/bin/python3 10
+    ;;
+  py3)
+    ;;
+  *)
+    usage
+esac
+
+case $2 in
+  cuda-included)
+    python build.py --enable_cuda --cudnn_path /usr/lib/x86_64-linux-gnu/
+    python build/include_cuda.py
+    ;;
+  cuda)
+    python build.py --enable_cuda --cudnn_path /usr/lib/x86_64-linux-gnu/
+    ;;
+  nocuda)
+    python build.py
+    ;;
+  *)
+    usage
+esac
+
+python setup.py bdist bdist_wheel
+cp -r dist/* /dist