Merge pull request #31 from Ericgig/bugfix.dep.version.update

Update for jax 0.4.23 and dimensions PR

.github/workflows/tests.yml

@@ -57,7 +57,7 @@ jobs:
         # to see if something hung.
         timeout-minutes: 60
         run: |
-          pytest --durations=0 --durations-min=1.0 --verbosity=1 --cov=qutip_jax --cov-report= --color=yes -W ignore::UserWarning:qutip
+          pytest --durations=0 --durations-min=1.0 --verbosity=1 --cov=qutip_jax --cov-report= --color=yes -W ignore::UserWarning:qutip -W "ignore:Complex dtype:UserWarning"
           # Above flags are:
           #  --durations=0 --durations-min=1.0
           #     at the end, show a list of all the tests that took longer than a

src/qutip_jax/jaxarray.py

@@ -1,6 +1,5 @@
 import jax.numpy as jnp
-from jax import tree_util
-from jax.config import config
+from jax import tree_util, config
 import numbers
 import numpy as np
 

src/qutip_jax/qutip_trees.py

@@ -9,9 +9,7 @@
 def qobj_tree_flatten(qobj):
     children = (qobj.to("jax").data,)
     aux_data = {
-        "dims": qobj.dims,
-        "type": qobj.type,
-        "superrep": qobj.superrep,
+        "_dims": qobj._dims,
         # Attribute that depend on the data are not safe to be set.
         "_isherm": None,
         "_isunitary": None,