fix qutip_jax import

qutip · Feb 28, 2024 · e952b44 · e952b44
1 parent fd7abcf
commit e952b44
Show file tree

Hide file tree

Showing 3 changed files with 4 additions and 0 deletions.
diff --git a/mypy.ini b/mypy.ini
@@ -5,5 +5,7 @@ ignore_missing_imports = True
 ignore_missing_imports = True
 [mypy-diffrax.*]
 ignore_missing_imports = True
+[mypy-qutip_jax.*]
+ignore_missing_imports = True
 [mypy-qutip_qtrl.*]
 ignore_missing_imports = True
diff --git a/src/qutip_qoc/joat.py b/src/qutip_qoc/joat.py
@@ -10,6 +10,7 @@
 import jax
 from jax import custom_jvp
 import jax.numpy as jnp
+import qutip_jax  # noqa: F401
 
 __all__ = ["JOAT", "Multi_JOAT"]
 

diff --git a/tests/test_analytical_pulses.py b/tests/test_analytical_pulses.py
@@ -4,6 +4,7 @@
 
 import pytest
 import qutip as qt
+import qutip_jax  # noqa: F401
 import numpy as np
 import jax.numpy as jnp
 import collections