From a45dc7616395709ce29c1ecb8fb8662679e923c6 Mon Sep 17 00:00:00 2001
From: Kaushik Kulkarni <kaushikcfd@gmail.com>
Date: Sun, 18 Jul 2021 22:45:35 -0500
Subject: [PATCH] CI: adds jax to requirements

---
 .gitlab-ci.yml          | 3 +++
 .test-conda-env-py3.yml | 1 +
 2 files changed, 4 insertions(+)

diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index a321fa4..00ebda7 100644
--- a/.gitlab-ci.yml
+++ b/.gitlab-ci.yml
@@ -2,6 +2,7 @@ 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
     . ./build-and-test-py-project.sh
   tags:
@@ -18,6 +19,7 @@ Python 3 Nvidia Titan V:
   script: |
     export PY_EXE=python3
     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
   tags:
@@ -74,6 +76,7 @@ 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
   tags:
diff --git a/.test-conda-env-py3.yml b/.test-conda-env-py3.yml
index 907bcb5..15b6496 100644
--- a/.test-conda-env-py3.yml
+++ b/.test-conda-env-py3.yml
@@ -13,3 +13,4 @@ dependencies:
 - pyopencl
 - islpy
 - pip
+- jax
-- 
GitLab