added jax install to workflow

.github/workflows/ci.yml

@@ -25,7 +25,7 @@ jobs:
       - name: Install dependencies 🛠
         run: |
           python -m pip install --upgrade pip setuptools
-          pip install flake8 codecov pytest-cov wheel matplotlib seaborn mne
+          pip install flake8 codecov pytest-cov wheel matplotlib seaborn mne jax
           if [ -f requirements.txt ]; then pip install -r requirements.txt; fi
           pip install .
       - name: Lint with flake8 🎓