added more mkdocs dependencies to ci and setup.py

.github/workflows/ci.yml

@@ -25,5 +25,5 @@ jobs:
           path: .cache
           restore-keys: |
             mkdocs-material-
-      - run: pip install mkdocs-material 
+      - run: pip install mkdocs mkdocs-material mkdocs-literate-nav mkdocs-section-index mkdocs-gen-files mkdocstrings-python
       - run: mkdocs gh-deploy --force

jaxquantum/quantum/base.py

@@ -78,7 +78,8 @@ def dag(op: jnp.ndarray) -> jnp.ndarray:
     Returns:
         conjugate transpose of op
     """
-    return jnp.conj(op).T
+    # return jnp.conj(op).T
+    return jnp.moveaxis(jnp.conj(op), -1, -2) # transposes last two axes, good for batching
 
 
 def ket2dm(ket: jnp.ndarray) -> jnp.ndarray:

setup.py

@@ -24,7 +24,12 @@
         "black"
     ],
     "docs": [
-        "mkdocs-material"
+        "mkdocs",
+        "mkdocs-material",
+        "mkdocs-literate-nav",
+        "mkdocs-section-index",
+        "mkdocs-gen-files",
+        "mkdocstrings-python",
     ]
 }