fix jax version bug

.github/workflows/test.yml

@@ -63,6 +63,11 @@ jobs:
         run: |
           python -m pip install 'git+https://github.com/qutip/qutip-qtrl.git'
 
+      - name: Install qutip-jax from GitHub
+        if: ${{ matrix.qutip-version == '@master'}}
+        run: |
+          python -m pip install 'git+https://github.com/qutip/qutip-jax.git'
+
       - name: Install qutip-qoc
         # Installing in-place so that coveralls can locate the source code.
         run: |

doc/requirements.txt

@@ -1,7 +1,7 @@
 numpy>=1.16.6
 scipy>=1.10.1
-jax==0.4.23
-jaxlib==0.4.23
+jax>=0.4.23
+jaxlib>=0.4.23
 qutip>=5.0.1
 cython>=0.29.33
 numpydoc==1.5.0

doc/rtd-environment.yml

@@ -5,7 +5,7 @@ dependencies:
   - numpy>=1.16.6
   - scipy>=1.10.1
   - jax>=0.4.23
-  - jaxlib==0.4.23
+  - jaxlib>=0.4.23
   - cython>=0.29.33
   - sphinx==6.1.3
   - sphinx_rtd_theme==1.2.0

requirements.txt

@@ -1,9 +1,9 @@
 cython>=1.0
 numpy>=1.16.6
 scipy>=1.10.1
-jax==0.4.23
-jaxlib==0.4.23
+jax>=0.4.23
+jaxlib>=0.4.23
 qutip>=5.0.1
 qutip-qtrl @ git+https://github.com/qutip/qutip-qtrl.git@master
-qutip-jax
+qutip-jax @ git+https://github.com/qutip/qutip-jax.git@master
 pre-commit

setup.cfg

@@ -33,7 +33,7 @@ install_requires =
     packaging
     qutip
     qutip-qtrl @ git+https://github.com/qutip/qutip-qtrl.git@master
-    qutip-jax
+    qutip-jax @ git+https://github.com/qutip/qutip-jax.git@master
 setup_requires =
     cython>=1.0
     packaging

src/qutip_qoc/_optimizer.py

@@ -206,7 +206,7 @@ def _global_local_optimization(
                     Sequence of ``(min, max)`` pairs for each element in
                     `guess`. None is used to specify no bound.
 
-    time_interval : :class:`qutip_qoc.TimeInterval`
+    time_interval : :class:`qutip_qoc._TimeInterval`
         Time interval for the optimization.
 
     time_options : dict, optional