Fix JAX eigh

Author: abhinavd

tensorcircuit/backends/jax_ops.py

@@ -151,11 +151,16 @@ def _QrGradSquareAndDeepMatrices(q: Array, r: Array, dq: Array, dr: Array) -> Ar
 
 @jax.custom_vjp
 def adaware_eigh(A: Array) -> Array:
-    return jnp.linalg.eigh(A)
+    result = jnp.linalg.eigh(A)
+    e = result.eigenvalues
+    v = result.eigenvectors
+    return e, v
 
 
 def jaxeigh_fwd(A: Array) -> Array:
-    e, v = jnp.linalg.eigh(A)
+    result = jnp.linalg.eigh(A)
+    e = result.eigenvalues
+    v = result.eigenvectors
     return (e, v), (A, e, v)