Avoid buggy version of JAX in CI

Related to jax-ml/jax#22751
pymc-devs · Aug 4, 2024 · 95db980 · 95db980
1 parent f9bf218
commit 95db980
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/.github/workflows/test.yml b/.github/workflows/test.yml
@@ -156,7 +156,7 @@ jobs:
 
           micromamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy pip mkl-service graphviz cython pytest coverage pytest-cov pytest-benchmark pytest-mock
           if [[ $INSTALL_NUMBA == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" "numba>=0.57"; fi
-          if [[ $INSTALL_JAX == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" jax jaxlib numpyro && pip install tensorflow-probability; fi
+          if [[ $INSTALL_JAX == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" "jax!=0.4.31" "jaxlib!=0.4.31" numpyro && pip install tensorflow-probability; fi
           if [[ $INSTALL_TORCH == "1" ]]; then micromamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" pytorch pytorch-cuda=12.1 -c pytorch -c nvidia; fi
           pip install pytest-sphinx