Revert "Pin jaxlib to a version <= 0.1.75"

.github/workflows/test.yml

@@ -143,7 +143,7 @@ jobs:
         run: |
           mamba install --yes -q "python~=${PYTHON_VERSION}=*_cpython" mkl numpy scipy pip mkl-service graphviz cython pytest coverage pytest-cov sympy
           if [[ $INSTALL_NUMBA == "1" ]]; then mamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" "numba>=0.55" numba-scipy; fi
-          mamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" jax "jaxlib <= 0.1.75"
+          mamba install --yes -q -c conda-forge "python~=${PYTHON_VERSION}=*_cpython" jax jaxlib
           pip install -e ./
           mamba list && pip freeze
           python -c 'import aesara; print(aesara.config.__str__(print_doc=False))'

requirements.txt

@@ -11,7 +11,7 @@ cython
 sympy
 versioneer
 jax
-jaxlib<=0.1.75
+jaxlib
 numba>=0.55
 numba-scipy>=0.3.0
 diff-cover