diff --git a/.gitlab-ci.yml b/.gitlab-ci.yml
index a321fa4489fbc9ae5336d0995ff892bc5aa35c52..00ebda72346f45ff23ef661ffc4de0154620e604 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 907bcb5315603d400b96fa5466a4fef2801954a4..15b64965a09c1ba884453e4ed9a25cfdc18f2511 100644
--- a/.test-conda-env-py3.yml
+++ b/.test-conda-env-py3.yml
@@ -13,3 +13,4 @@ dependencies:
 - pyopencl
 - islpy
 - pip
+- jax