Merge pull request jax-ml#4847 from hawkinsp:py39

PiperOrigin-RevId: 341494288
bileschi · Nov 9, 2020 · a3a55a9 · a3a55a9
2 parents 307b528 + 2c6f932
commit a3a55a9
Show file tree

Hide file tree

Showing 5 changed files with 7 additions and 4 deletions.
diff --git a/WORKSPACE b/WORKSPACE
@@ -28,10 +28,10 @@ http_archive(
 #    and update the sha256 with the result.
 http_archive(
     name = "org_tensorflow",
-    sha256 = "df332b5f174d9e759dbbad95d6fbdebe4d64d95cf547abdd0b7e1fea4398760c",
-    strip_prefix = "tensorflow-55a311cb735689a431c6aa9a6c765c5c5c034ede",
+    sha256 = "34f473fc2dc80ad375a196bf6e98812c65d2eeb39aed6521aa0b2b03d34eb520",
+    strip_prefix = "tensorflow-4cf039a9877a5e7a73cf3bc189af0bd03d2cb422",
     urls = [
-        "https://github.com/tensorflow/tensorflow/archive/55a311cb735689a431c6aa9a6c765c5c5c034ede.tar.gz",
+        "https://github.com/tensorflow/tensorflow/archive/4cf039a9877a5e7a73cf3bc189af0bd03d2cb422.tar.gz",
     ],
 )
 

diff --git a/build/Dockerfile b/build/Dockerfile
@@ -16,6 +16,7 @@ RUN /pyenv/bin/pyenv install 3.8.0
 RUN eval "$(/pyenv/bin/pyenv init -)" && /pyenv/bin/pyenv local 3.6.8 && pip install numpy==1.15.4 scipy cython setuptools wheel packaging six auditwheel
 RUN eval "$(/pyenv/bin/pyenv init -)" && /pyenv/bin/pyenv local 3.7.2 && pip install numpy==1.15.4 scipy cython setuptools wheel packaging six auditwheel
 RUN eval "$(/pyenv/bin/pyenv init -)" && /pyenv/bin/pyenv local 3.8.0 && pip install numpy==1.17.3 scipy cython setuptools wheel packaging six auditwheel
+RUN eval "$(/pyenv/bin/pyenv init -)" && /pyenv/bin/pyenv local 3.9.0 && pip install numpy==1.19.4 scipy==1.5.4 cython setuptools wheel packaging six auditwheel
 
 # Change the CUDA version if it doesn't match the installed version.
 ARG JAX_CUDA_VERSION=10.0

diff --git a/build/build.py b/build/build.py
@@ -179,6 +179,7 @@ def check_bazel_version(bazel_path, min_version, max_version):
 common --experimental_repo_remote_exec
 
 build --repo_env PYTHON_BIN_PATH="{python_bin_path}"
+build --action_env=PYENV_ROOT
 build --python_path="{python_bin_path}"
 build --repo_env TF_NEED_CUDA="{tf_need_cuda}"
 build --action_env TF_CUDA_COMPUTE_CAPABILITIES="{cuda_compute_capabilities}"

diff --git a/build/build_jaxlib_wheels.sh b/build/build_jaxlib_wheels.sh
@@ -1,7 +1,7 @@
 #!/bin/bash
 set -xev
 
-PYTHON_VERSIONS="3.6.8 3.7.2 3.8.0"
+PYTHON_VERSIONS="3.6.8 3.7.2 3.8.0 3.9.0"
 CUDA_VERSIONS="10.0 10.1 10.2 11.0 11.1"
 CUDA_VARIANTS="cuda" # "cuda-included"
 

diff --git a/build/build_jaxlib_wheels_macos.sh b/build/build_jaxlib_wheels_macos.sh
@@ -54,3 +54,4 @@ rm -fr build/dist
 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.2
+build_jax 3.9.0 cp39 1.19.4 1.5.4