Update for latest jax

qutip · Feb 12, 2024 · 6c1e52c · 6c1e52c
1 parent 8781aff
commit 6c1e52c
Showing 1 changed file with 1 addition and 2 deletions.
diff --git a/src/qutip_jax/jaxdia.py b/src/qutip_jax/jaxdia.py
@@ -1,7 +1,6 @@
 import jax.numpy as jnp
 import numpy as np
-from jax import tree_util, jit
-from jax.config import config
+from jax import tree_util, jit, config
 from qutip.core.data.extract import extract
 import qutip.core.data as _data
 import numpy as np