From f26c4ee5212fa1a76478b8ca825db7d097e40ad3 Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Thu, 12 May 2022 19:42:27 -0500
Subject: [PATCH 1/2] Fix Jax installation on Gitlab CI

---
 .gitlab-ci.yml | 9 ++++++---
 1 file changed, 6 insertions(+), 3 deletions(-)

diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index 00ebda72..7c9358a6 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -18,10 +18,12 @@ Python 3 POCL:
 Python 3 Nvidia Titan V:
   script: |
     export PY_EXE=python3
+    curl -L -O https://tiker.net/ci-support-v0
     export PYOPENCL_TEST=nvi:titan
-    export EXTRA_INSTALL="jax[cuda]"
-    curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project.sh
-    . ./build-and-test-py-project.sh
+    build_py_project_in_venv
+    pip install --upgrade "jax[cuda]" -f https://storage.googleapis.com/jax-releases/jax_releases.html
+    test_py_project
+
   tags:
   - python3
   - nvidia-titan-v
@@ -86,6 +88,7 @@ Pylint:
 
 Mypy:
   script: |
+    EXTRA_INSTALL="jax[cpu]"
     curl -L -O https://tiker.net/ci-support-v0
     . ./ci-support-v0
     build_py_project_in_venv

From 846c534f101d29f1b6b817823f684b0de254a11c Mon Sep 17 00:00:00 2001
From: Andreas Kloeckner <inform@tiker.net>
Date: Thu, 12 May 2022 19:43:07 -0500
Subject: [PATCH 2/2] Stop setting PY_EXE in Gitlab CI config

---
 .gitlab-ci.yml | 4 ----
 1 file changed, 4 deletions(-)

diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index 7c9358a6..be4962cc 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -1,6 +1,5 @@
 Python 3 POCL:
   script: |
-    export PY_EXE=python3
     export PYOPENCL_TEST=portable:pthread
     export EXTRA_INSTALL="jax[cpu]"
     curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-and-test-py-project.sh
@@ -17,7 +16,6 @@ Python 3 POCL:
 
 Python 3 Nvidia Titan V:
   script: |
-    export PY_EXE=python3
     curl -L -O https://tiker.net/ci-support-v0
     export PYOPENCL_TEST=nvi:titan
     build_py_project_in_venv
@@ -36,7 +34,6 @@ Python 3 Nvidia Titan V:
 Python 3 POCL Examples:
   script:
   - test -n "$SKIP_EXAMPLES" && exit
-  - export PY_EXE=python3
   - export PYOPENCL_TEST=portable:pthread
   - curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/main/build-py-project-and-run-examples.sh
   - ". ./build-py-project-and-run-examples.sh"
@@ -77,7 +74,6 @@ Flake8:
 
 Pylint:
   script: |
-    export PY_EXE=python3
     EXTRA_INSTALL="jax[cpu]"
     curl -L -O https://gitlab.tiker.net/inducer/ci-support/raw/master/prepare-and-run-pylint.sh
     . ./prepare-and-run-pylint.sh "$CI_PROJECT_NAME" examples/*.py test/test_*.py