Merge pull request #92 from Ericgig/v0.1.1

V0.1.1

Author: Ericgig

VERSION

@@ -1 +1 @@
-0.1.0
+0.1.1

src/qutip_jax/qobjevo.py

@@ -14,6 +14,14 @@
 from functools import partial
 
 
+try:
+    # Pre jac 0.6.1
+    PjitFunction = jaxlib.xla_extension.PjitFunction
+except AttributeError:
+    # Post jac 0.6.1
+    PjitFunction = jaxlib._jax.PjitFunction
+
+
 __all__ = []
 
 
@@ -146,7 +154,7 @@ def unflatten(cls, aux_data, children):
         )
 
 
-coefficient_builders[jaxlib.xla_extension.PjitFunction] = JaxJitCoeff
+coefficient_builders[PjitFunction] = JaxJitCoeff
 jax.tree_util.register_pytree_node(
     JaxJitCoeff, JaxJitCoeff.flatten, JaxJitCoeff.unflatten
 )

tests/test_unary.py

@@ -60,7 +60,7 @@ def _inv_jax(matrix):
         data.add(
             matrix,
             data.diag(
-                [1.1] * matrix.shape[0], shape=matrix.shape, dtype="JaxArray"
+                [2.0] * matrix.shape[0], shape=matrix.shape, dtype="JaxArray"
             ),
         )
     )