ci: pin jax<0.4.31 for mypy on gitlab

Author: alexfikl

.gitlab-ci.yml

@@ -107,11 +107,14 @@ Pylint:
 
 Mypy:
   script: |
-    EXTRA_INSTALL="jax[cpu]"
+    # NOTE: jax>=0.4.31 requires python 3.10 and uses pattern matching
+    # which conflicts with our mypy.python_version = '3.8' setting
+    EXTRA_INSTALL="mypy pytest jax[cpu]<0.4.31"
+
     curl -L -O https://tiker.net/ci-support-v0
     . ./ci-support-v0
+
     build_py_project_in_venv
-    python -m pip install mypy pytest
     ./run-mypy.sh
   tags:
   - python3

pyproject.toml

@@ -119,6 +119,7 @@ known-local-folder = [
 lines-after-imports = 2
 
 [tool.mypy]
+# TODO: unpin jax version on CI when this gets bumped to 3.10
 python_version = "3.8"
 warn_unused_ignores = true
 # TODO: enable this