[ci] Use Jax <= 0.3.15
Newer versions fail to import stax
@@ -26,7 +26,7 @@ jobs:
- name: Install dependencies
run: |
python -m pip install --upgrade pip
- pip install --quiet tensorflow tensorflow-probability torch jax jaxlib plotly nbformat ipython pylint coverage pytest
+ pip install --quiet tensorflow tensorflow-probability torch "jax<=0.3.15" "jaxlib<=0.3.15" plotly nbformat ipython pylint coverage pytest
pip install .
- name: Test with pytest
