ci: pin jax<0.4.31 for mypy

Author: alexfikl

.github/workflows/ci.yml

@@ -43,14 +43,16 @@ jobs:
         runs-on: ubuntu-latest
         steps:
         -   uses: actions/checkout@v4
-        -
-            uses: actions/setup-python@v5
-            with:
-                python-version: '3.x'
         -   name: "Main Script"
             run: |
                 curl -L -O https://tiker.net/ci-support-v0
                 . ./ci-support-v0
+
+                # NOTE: jax>=0.4.31 requires python 3.10 and uses pattern matching
+                # which conflicts with our mypy.python_version = '3.8' setting
+                CONDA_ENVIRONMENT=.test-conda-env-py3.yml
+                sed -i "/jax/jax<0.4.31/" "$CONDA_ENVIRONMENT"
+
                 build_py_project_in_conda_env
                 python -m pip install mypy pytest
                 ./run-mypy.sh

.test-conda-env-py3.yml

@@ -14,4 +14,4 @@ dependencies:
 - pyopencl
 - islpy
 - pip
-- jax
+- jax<0.4.31