From e952b4415368bb49ef49b18bffa51bba4c9a522e Mon Sep 17 00:00:00 2001
From: Patrick <>
Date: Wed, 28 Feb 2024 09:51:46 +0100
Subject: [PATCH] fix qutip_jax import

---
 mypy.ini                        | 2 ++
 src/qutip_qoc/joat.py           | 1 +
 tests/test_analytical_pulses.py | 1 +
 3 files changed, 4 insertions(+)

diff --git a/mypy.ini b/mypy.ini
index 69fadbd..84acb73 100644
--- 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
index e39f347..38e39be 100644
--- 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
index b348f38..1c515f7 100644
--- 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