setting jax version warning to only trigger if the version is between…

… 0.4.4 and 0.4.6 and the os flag hasn't been set

qiskit_dynamics/dispatch/backends/jax.py

@@ -25,12 +25,14 @@
     import warnings
 
     if version.parse(jax.__version__) >= version.parse("0.4.4"):
-        warnings.warn(
-            "The functionality in the perturbation module of Qiskit Dynamics requires a JAX "
-            "version <= 0.4.6, due to a bug in JAX versions > 0.4.6. For versions 0.4.4, 0.4.5, "
-            "and 0.4.6, using the perturbation module functionality requires setting "
-            "os.environ['JAX_JIT_PJIT_API_MERGE'] = '0' before importing JAX."
-        )
+        import os
+        if version.parse(jax.__version__) > version.parse("0.4.6") or os.environ.get('JAX_JIT_PJIT_API_MERGE', None) != '0':
+            warnings.warn(
+                "The functionality in the perturbation module of Qiskit Dynamics requires a JAX "
+                "version <= 0.4.6, due to a bug in JAX versions > 0.4.6. For versions 0.4.4, "
+                "0.4.5, and 0.4.6, using the perturbation module functionality requires setting "
+                "os.environ['JAX_JIT_PJIT_API_MERGE'] = '0' before importing JAX or Dynamics."
+            )
 
     JAX_TYPES = (Array, Tracer)
 

tox.ini

@@ -41,8 +41,8 @@ commands = black {posargs} qiskit_dynamics test
 usedevelop = False
 deps =
     -r{toxinidir}/requirements-dev.txt
-    jax<=0.4.3
-    jaxlib<=0.4.3
+    jax<=0.4.6
+    jaxlib<=0.4.6
     diffrax
 commands =
   sphinx-build -j auto -b html -W {posargs} docs/ docs/_build/html