diff --git a/docs/notebooks/Quick_Start.ipynb b/docs/notebooks/Quick_Start.ipynb
index dbeca387..150f5c13 100644
--- a/docs/notebooks/Quick_Start.ipynb
+++ b/docs/notebooks/Quick_Start.ipynb
@@ -27,7 +27,7 @@
     {
      "data": {
       "text/plain": [
-       "<torch._C.Generator at 0x7fe235dde670>"
+       "<torch._C.Generator at 0x7f0f65a5a3f0>"
       ]
      },
      "execution_count": 1,
@@ -38,6 +38,10 @@
    "source": [
     "import cola\n",
     "import torch\n",
+    "import numpy as np\n",
+    "import warnings\n",
+    "\n",
+    "warnings.filterwarnings(\"ignore\")\n",
     "torch.manual_seed(21)"
    ]
   },
@@ -196,11 +200,10 @@
     }
    ],
    "source": [
-    "import numpy as np\n",
     "perm = torch.randperm(200)\n",
     "P = cola.ops.Permutation(perm)\n",
     "v = torch.randn(200)\n",
-    "print((P@v)[:5],v[perm[:5]])"
+    "print((P @ v)[:5], v[perm[:5]])"
    ]
   },
   {
@@ -246,8 +249,8 @@
     }
    ],
    "source": [
-    "D = cola.diag(torch.tensor([1,2,3,4.]))\n",
-    "DT = T[:4,:4]-D\n",
+    "D = cola.diag(torch.tensor([1, 2, 3, 4.]))\n",
+    "DT = T[:4, :4] - D\n",
     "print(DT.to_dense())"
    ]
   },
@@ -294,7 +297,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 9,
    "id": "06defdf2",
    "metadata": {},
    "outputs": [
@@ -312,13 +315,13 @@
     }
    ],
    "source": [
-    "K = cola.kron(A,D)\n",
+    "K = cola.kron(A, D)\n",
     "print(f\"{K[:5,:5].to_dense()}\\n shape: {K.shape}, object: {K}\")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 10,
    "id": "f660e95b",
    "metadata": {},
    "outputs": [
@@ -328,14 +331,14 @@
        "<2100x2100 BlockDiag[cola.ops.operators.Kronecker[cola.ops.operators.Dense, cola.ops.operators.Permutation], cola.ops.operators.Tridiagonal, cola.ops.operators.Diagonal] with dtype=torch.float32>"
       ]
      },
-     "execution_count": 11,
+     "execution_count": 10,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "d = 1/(1+torch.linspace(.2,10,1000))\n",
-    "X = cola.block_diag(cola.kron(A,P), T, cola.diag(d))\n",
+    "d = 1 / (1 + torch.linspace(.2, 10, 1000))\n",
+    "X = cola.block_diag(cola.kron(A, P), T, cola.diag(d))\n",
     "X"
    ]
   },
@@ -365,7 +368,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 24,
+   "execution_count": 11,
    "id": "4e803b0c",
    "metadata": {},
    "outputs": [
@@ -375,13 +378,13 @@
        "torch.Size([2100, 5])"
       ]
      },
-     "execution_count": 24,
+     "execution_count": 11,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "X[:,:5].to_dense().shape"
+    "X[:, :5].to_dense().shape"
    ]
   },
   {
@@ -425,7 +428,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": 14,
    "id": "04d88f0f",
    "metadata": {},
    "outputs": [
@@ -444,7 +447,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 44,
+   "execution_count": 15,
    "id": "b6432c5b",
    "metadata": {},
    "outputs": [
@@ -457,11 +460,11 @@
       "        [0.0932, 0.2152, 0.3083, 0.2139, 0.0864],\n",
       "        [0.0288, 0.0930, 0.2139, 0.3016, 0.1865],\n",
       "        [0.0066, 0.0275, 0.0864, 0.1865, 0.2153]])\n",
-      "scipy expm: [[0.30850828 0.21526888 0.09323523 0.02876081 0.0066488 ]\n",
-      " [0.21526891 0.30850452 0.2152388  0.09302244 0.02746145]\n",
-      " [0.09323531 0.2152388  0.30829173 0.21393944 0.08637362]\n",
-      " [0.02876082 0.09302243 0.2139395  0.30164295 0.18647799]\n",
-      " [0.0066488  0.02746145 0.08637366 0.18647805 0.21526927]]\n"
+      "scipy expm: [[0.30850834 0.21526882 0.09323522 0.02876082 0.0066488 ]\n",
+      " [0.21526891 0.30850452 0.21523876 0.09302247 0.02746145]\n",
+      " [0.09323531 0.21523881 0.30829167 0.21393949 0.08637363]\n",
+      " [0.02876082 0.09302244 0.21393946 0.30164295 0.18647799]\n",
+      " [0.0066488  0.02746145 0.08637364 0.18647805 0.21526927]]\n"
      ]
     }
    ],
@@ -473,42 +476,47 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 36,
+   "execution_count": 16,
    "id": "23e4504e",
    "metadata": {},
    "outputs": [
     {
-     "data": {
-      "text/plain": [
-       "tensor(0.+0.j)"
-      ]
-     },
-     "execution_count": 36,
-     "metadata": {},
-     "output_type": "execute_result"
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "CoLA exp(X): tensor([[0.3085, 0.2153, 0.0932, 0.0288, 0.0066],\n",
+      "        [0.2153, 0.3085, 0.2152, 0.0930, 0.0275],\n",
+      "        [0.0932, 0.2152, 0.3083, 0.2139, 0.0864],\n",
+      "        [0.0288, 0.0930, 0.2139, 0.3016, 0.1865],\n",
+      "        [0.0066, 0.0275, 0.0864, 0.1865, 0.2153]])\n",
+      "scipy expm: [[0.30850834 0.21526882 0.09323522 0.02876082 0.0066488 ]\n",
+      " [0.21526891 0.30850452 0.21523876 0.09302247 0.02746145]\n",
+      " [0.09323531 0.21523881 0.30829167 0.21393949 0.08637363]\n",
+      " [0.02876082 0.09302244 0.21393946 0.30164295 0.18647799]\n",
+      " [0.0066488  0.02746145 0.08637364 0.18647805 0.21526927]]\n"
+     ]
     }
    ],
    "source": [
-    "print(f\"CoLA exp(X): {cola.exp(T,method='iterative')[-5:,-5:].to_dense().real}\\nscipy expm: {scipy_expm[-5:,-5:]}\")\n",
-    "#TODO Andres: investigate why iterative exp not working for big matrices"
+    "print(f\"CoLA exp(X): {cola.exp(T, method='iterative')[-5:,-5:].to_dense().real}\\nscipy expm: {scipy_expm[-5:,-5:]}\")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 54,
+   "execution_count": 17,
    "id": "32ad8a12",
    "metadata": {},
    "outputs": [],
    "source": [
-    "C = X.T@X\n",
+    "C = X.T @ X\n",
     "# let's annotate this matrix as positive definite to speed up the computation\n",
     "C = cola.PSD(C)\n",
-    "e,v = cola.eig(C,eig_slice=np.s_[-3:],max_iters=100)"
+    "e, v = cola.eig(C, eig_slice=np.s_[-3:], max_iters=100)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 55,
+   "execution_count": 18,
    "id": "2d0f3297",
    "metadata": {},
    "outputs": [
@@ -516,7 +524,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "tensor([15.9769, 15.9922, 15.9984]) tensor([15.9922, 15.9965, 15.9991])\n"
+      "tensor([15.9803, 15.9892, 15.9981]) tensor([15.9922, 15.9965, 15.9991])\n"
      ]
     }
    ],
@@ -526,47 +534,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 57,
+   "execution_count": 19,
    "id": "477eae7e",
    "metadata": {},
    "outputs": [
     {
-     "ename": "KeyboardInterrupt",
-     "evalue": "",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[57], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m e,v \u001b[39m=\u001b[39m cola\u001b[39m.\u001b[39;49meig(C,eig_slice\u001b[39m=\u001b[39;49mnp\u001b[39m.\u001b[39;49ms_[:\u001b[39m3\u001b[39;49m])\n\u001b[1;32m      2\u001b[0m \u001b[39mprint\u001b[39m(e, torch\u001b[39m.\u001b[39mlinalg\u001b[39m.\u001b[39meigh(C\u001b[39m.\u001b[39mto_dense())[\u001b[39m0\u001b[39m][:\u001b[39m3\u001b[39m])\n",
-      "File \u001b[0;32m~/anaconda3/envs/cola/lib/python3.10/site-packages/plum/function.py:438\u001b[0m, in \u001b[0;36mFunction.__call__\u001b[0;34m(self, *args, **kw_args)\u001b[0m\n\u001b[1;32m    436\u001b[0m     method, return_type, loginfo \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mresolve_method(args, types)\n\u001b[1;32m    437\u001b[0m logging\u001b[39m.\u001b[39minfo(\u001b[39m\"\u001b[39m\u001b[39m%s\u001b[39;00m\u001b[39m\"\u001b[39m,loginfo)\n\u001b[0;32m--> 438\u001b[0m \u001b[39mreturn\u001b[39;00m _convert(method(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkw_args), return_type)\n",
-      "File \u001b[0;32m~/cola/cola/linalg/eigs.py:74\u001b[0m, in \u001b[0;36meig\u001b[0;34m(A, **kwargs)\u001b[0m\n\u001b[1;32m     72\u001b[0m     \u001b[39mreturn\u001b[39;00m eig_vals[eig_slice], Unitary(lazify(eig_vecs[:, eig_slice]))\n\u001b[1;32m     73\u001b[0m \u001b[39melif\u001b[39;00m method \u001b[39min\u001b[39;00m (\u001b[39m'\u001b[39m\u001b[39mlanczos\u001b[39m\u001b[39m'\u001b[39m, \u001b[39m'\u001b[39m\u001b[39miterative\u001b[39m\u001b[39m'\u001b[39m) \u001b[39mor\u001b[39;00m (method \u001b[39m==\u001b[39m \u001b[39m'\u001b[39m\u001b[39mauto\u001b[39m\u001b[39m'\u001b[39m \u001b[39mand\u001b[39;00m prod(A\u001b[39m.\u001b[39mshape) \u001b[39m>\u001b[39m\u001b[39m=\u001b[39m \u001b[39m1e6\u001b[39m):\n\u001b[0;32m---> 74\u001b[0m     eig_vals, eig_vecs \u001b[39m=\u001b[39m eig(LanczosDecomposition(A, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkws), eig_slice\u001b[39m=\u001b[39;49meig_slice)\n\u001b[1;32m     75\u001b[0m     \u001b[39mreturn\u001b[39;00m eig_vals, eig_vecs\n\u001b[1;32m     76\u001b[0m \u001b[39melse\u001b[39;00m:\n",
-      "File \u001b[0;32m~/anaconda3/envs/cola/lib/python3.10/site-packages/plum/function.py:438\u001b[0m, in \u001b[0;36mFunction.__call__\u001b[0;34m(self, *args, **kw_args)\u001b[0m\n\u001b[1;32m    436\u001b[0m     method, return_type, loginfo \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mresolve_method(args, types)\n\u001b[1;32m    437\u001b[0m logging\u001b[39m.\u001b[39minfo(\u001b[39m\"\u001b[39m\u001b[39m%s\u001b[39;00m\u001b[39m\"\u001b[39m,loginfo)\n\u001b[0;32m--> 438\u001b[0m \u001b[39mreturn\u001b[39;00m _convert(method(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkw_args), return_type)\n",
-      "File \u001b[0;32m~/cola/cola/decompositions.py:51\u001b[0m, in \u001b[0;36meig\u001b[0;34m(QH, **kwargs)\u001b[0m\n\u001b[1;32m     48\u001b[0m \u001b[39m@eig\u001b[39m\u001b[39m.\u001b[39mdispatch\n\u001b[1;32m     49\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39meig\u001b[39m(QH: UnitaryDecomposition, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs):\n\u001b[1;32m     50\u001b[0m     Q, H, xnp \u001b[39m=\u001b[39m QH\u001b[39m.\u001b[39mQ, QH\u001b[39m.\u001b[39mM, QH\u001b[39m.\u001b[39mxnp\n\u001b[0;32m---> 51\u001b[0m     eig_vals, eig_vecs \u001b[39m=\u001b[39m eig(H, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m     52\u001b[0m     eig_vecs \u001b[39m=\u001b[39m xnp\u001b[39m.\u001b[39mcast(Q\u001b[39m.\u001b[39mto_dense(), dtype\u001b[39m=\u001b[39meig_vecs\u001b[39m.\u001b[39mdtype) \u001b[39m@\u001b[39m eig_vecs\u001b[39m.\u001b[39mto_dense()\n\u001b[1;32m     53\u001b[0m     eig_vecs \u001b[39m=\u001b[39m Unitary(lazify(eig_vecs))\n",
-      "    \u001b[0;31m[... skipping similar frames: Function.__call__ at line 438 (1 times)]\u001b[0m\n",
-      "File \u001b[0;32m~/cola/cola/linalg/eigs.py:74\u001b[0m, in \u001b[0;36meig\u001b[0;34m(A, **kwargs)\u001b[0m\n\u001b[1;32m     72\u001b[0m     \u001b[39mreturn\u001b[39;00m eig_vals[eig_slice], Unitary(lazify(eig_vecs[:, eig_slice]))\n\u001b[1;32m     73\u001b[0m \u001b[39melif\u001b[39;00m method \u001b[39min\u001b[39;00m (\u001b[39m'\u001b[39m\u001b[39mlanczos\u001b[39m\u001b[39m'\u001b[39m, \u001b[39m'\u001b[39m\u001b[39miterative\u001b[39m\u001b[39m'\u001b[39m) \u001b[39mor\u001b[39;00m (method \u001b[39m==\u001b[39m \u001b[39m'\u001b[39m\u001b[39mauto\u001b[39m\u001b[39m'\u001b[39m \u001b[39mand\u001b[39;00m prod(A\u001b[39m.\u001b[39mshape) \u001b[39m>\u001b[39m\u001b[39m=\u001b[39m \u001b[39m1e6\u001b[39m):\n\u001b[0;32m---> 74\u001b[0m     eig_vals, eig_vecs \u001b[39m=\u001b[39m eig(LanczosDecomposition(A, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkws), eig_slice\u001b[39m=\u001b[39;49meig_slice)\n\u001b[1;32m     75\u001b[0m     \u001b[39mreturn\u001b[39;00m eig_vals, eig_vecs\n\u001b[1;32m     76\u001b[0m \u001b[39melse\u001b[39;00m:\n",
-      "    \u001b[0;31m[... skipping similar frames: Function.__call__ at line 438 (1 times)]\u001b[0m\n",
-      "File \u001b[0;32m~/cola/cola/decompositions.py:51\u001b[0m, in \u001b[0;36meig\u001b[0;34m(QH, **kwargs)\u001b[0m\n\u001b[1;32m     48\u001b[0m \u001b[39m@eig\u001b[39m\u001b[39m.\u001b[39mdispatch\n\u001b[1;32m     49\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39meig\u001b[39m(QH: UnitaryDecomposition, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs):\n\u001b[1;32m     50\u001b[0m     Q, H, xnp \u001b[39m=\u001b[39m QH\u001b[39m.\u001b[39mQ, QH\u001b[39m.\u001b[39mM, QH\u001b[39m.\u001b[39mxnp\n\u001b[0;32m---> 51\u001b[0m     eig_vals, eig_vecs \u001b[39m=\u001b[39m eig(H, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m     52\u001b[0m     eig_vecs \u001b[39m=\u001b[39m xnp\u001b[39m.\u001b[39mcast(Q\u001b[39m.\u001b[39mto_dense(), dtype\u001b[39m=\u001b[39meig_vecs\u001b[39m.\u001b[39mdtype) \u001b[39m@\u001b[39m eig_vecs\u001b[39m.\u001b[39mto_dense()\n\u001b[1;32m     53\u001b[0m     eig_vecs \u001b[39m=\u001b[39m Unitary(lazify(eig_vecs))\n",
-      "    \u001b[0;31m[... skipping similar frames: Function.__call__ at line 438 (156 times), eig at line 74 (78 times), eig at line 51 (77 times)]\u001b[0m\n",
-      "File \u001b[0;32m~/cola/cola/decompositions.py:51\u001b[0m, in \u001b[0;36meig\u001b[0;34m(QH, **kwargs)\u001b[0m\n\u001b[1;32m     48\u001b[0m \u001b[39m@eig\u001b[39m\u001b[39m.\u001b[39mdispatch\n\u001b[1;32m     49\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39meig\u001b[39m(QH: UnitaryDecomposition, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs):\n\u001b[1;32m     50\u001b[0m     Q, H, xnp \u001b[39m=\u001b[39m QH\u001b[39m.\u001b[39mQ, QH\u001b[39m.\u001b[39mM, QH\u001b[39m.\u001b[39mxnp\n\u001b[0;32m---> 51\u001b[0m     eig_vals, eig_vecs \u001b[39m=\u001b[39m eig(H, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m     52\u001b[0m     eig_vecs \u001b[39m=\u001b[39m xnp\u001b[39m.\u001b[39mcast(Q\u001b[39m.\u001b[39mto_dense(), dtype\u001b[39m=\u001b[39meig_vecs\u001b[39m.\u001b[39mdtype) \u001b[39m@\u001b[39m eig_vecs\u001b[39m.\u001b[39mto_dense()\n\u001b[1;32m     53\u001b[0m     eig_vecs \u001b[39m=\u001b[39m Unitary(lazify(eig_vecs))\n",
-      "File \u001b[0;32m~/anaconda3/envs/cola/lib/python3.10/site-packages/plum/function.py:438\u001b[0m, in \u001b[0;36mFunction.__call__\u001b[0;34m(self, *args, **kw_args)\u001b[0m\n\u001b[1;32m    436\u001b[0m     method, return_type, loginfo \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mresolve_method(args, types)\n\u001b[1;32m    437\u001b[0m logging\u001b[39m.\u001b[39minfo(\u001b[39m\"\u001b[39m\u001b[39m%s\u001b[39;00m\u001b[39m\"\u001b[39m,loginfo)\n\u001b[0;32m--> 438\u001b[0m \u001b[39mreturn\u001b[39;00m _convert(method(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkw_args), return_type)\n",
-      "File \u001b[0;32m~/cola/cola/linalg/eigs.py:74\u001b[0m, in \u001b[0;36meig\u001b[0;34m(A, **kwargs)\u001b[0m\n\u001b[1;32m     72\u001b[0m     \u001b[39mreturn\u001b[39;00m eig_vals[eig_slice], Unitary(lazify(eig_vecs[:, eig_slice]))\n\u001b[1;32m     73\u001b[0m \u001b[39melif\u001b[39;00m method \u001b[39min\u001b[39;00m (\u001b[39m'\u001b[39m\u001b[39mlanczos\u001b[39m\u001b[39m'\u001b[39m, \u001b[39m'\u001b[39m\u001b[39miterative\u001b[39m\u001b[39m'\u001b[39m) \u001b[39mor\u001b[39;00m (method \u001b[39m==\u001b[39m \u001b[39m'\u001b[39m\u001b[39mauto\u001b[39m\u001b[39m'\u001b[39m \u001b[39mand\u001b[39;00m prod(A\u001b[39m.\u001b[39mshape) \u001b[39m>\u001b[39m\u001b[39m=\u001b[39m \u001b[39m1e6\u001b[39m):\n\u001b[0;32m---> 74\u001b[0m     eig_vals, eig_vecs \u001b[39m=\u001b[39m eig(LanczosDecomposition(A, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkws), eig_slice\u001b[39m=\u001b[39meig_slice)\n\u001b[1;32m     75\u001b[0m     \u001b[39mreturn\u001b[39;00m eig_vals, eig_vecs\n\u001b[1;32m     76\u001b[0m \u001b[39melse\u001b[39;00m:\n",
-      "File \u001b[0;32m~/cola/cola/algorithms/lanczos.py:103\u001b[0m, in \u001b[0;36mLanczosDecomposition\u001b[0;34m(A, start_vector, max_iters, tol, pbar)\u001b[0m\n\u001b[1;32m    100\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mLanczosDecomposition\u001b[39m(A: LinearOperator, start_vector\u001b[39m=\u001b[39m\u001b[39mNone\u001b[39;00m, max_iters\u001b[39m=\u001b[39m\u001b[39m100\u001b[39m, tol\u001b[39m=\u001b[39m\u001b[39m1e-7\u001b[39m, pbar\u001b[39m=\u001b[39m\u001b[39mFalse\u001b[39;00m):\n\u001b[1;32m    101\u001b[0m \u001b[39m    \u001b[39m\u001b[39m\"\"\" Provides the Lanczos decomposition of a matrix A = Q T Q^*.\u001b[39;00m\n\u001b[1;32m    102\u001b[0m \u001b[39m    LinearOperator form of lanczos, see lanczos for arguments.\"\"\"\u001b[39;00m\n\u001b[0;32m--> 103\u001b[0m     Q, T, info \u001b[39m=\u001b[39m lanczos_decomp(A\u001b[39m=\u001b[39;49mA, start_vector\u001b[39m=\u001b[39;49mstart_vector, max_iters\u001b[39m=\u001b[39;49mmax_iters, tol\u001b[39m=\u001b[39;49mtol,\n\u001b[1;32m    104\u001b[0m                                 pbar\u001b[39m=\u001b[39;49mpbar)\n\u001b[1;32m    105\u001b[0m     A_approx \u001b[39m=\u001b[39m cola\u001b[39m.\u001b[39mUnitaryDecomposition(lazify(Q), SelfAdjoint(lazify(T)))\n\u001b[1;32m    106\u001b[0m     A_approx\u001b[39m.\u001b[39minfo \u001b[39m=\u001b[39m info\n",
-      "File \u001b[0;32m~/cola/cola/algorithms/lanczos.py:132\u001b[0m, in \u001b[0;36mlanczos_decomp\u001b[0;34m(A, start_vector, max_iters, tol, pbar)\u001b[0m\n\u001b[1;32m    130\u001b[0m \u001b[39mif\u001b[39;00m start_vector \u001b[39mis\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m    131\u001b[0m     start_vector \u001b[39m=\u001b[39m xnp\u001b[39m.\u001b[39mfixed_normal_samples((A\u001b[39m.\u001b[39mshape[\u001b[39m0\u001b[39m], \u001b[39m1\u001b[39m))\n\u001b[0;32m--> 132\u001b[0m alpha, beta, vec, iters, info \u001b[39m=\u001b[39m lanczos_parts(A\u001b[39m=\u001b[39;49mA, rhs\u001b[39m=\u001b[39;49mstart_vector, max_iters\u001b[39m=\u001b[39;49mmax_iters,\n\u001b[1;32m    133\u001b[0m                                               tol\u001b[39m=\u001b[39;49mtol, pbar\u001b[39m=\u001b[39;49mpbar)\n\u001b[1;32m    134\u001b[0m alpha, beta, Q \u001b[39m=\u001b[39m alpha[\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m, :iters \u001b[39m-\u001b[39m \u001b[39m1\u001b[39m], beta[\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m, :iters], vec[\u001b[39m0\u001b[39m, :, :iters]\n\u001b[1;32m    135\u001b[0m T \u001b[39m=\u001b[39m construct_tridiagonal_batched(alpha, beta, alpha)[\u001b[39m0\u001b[39m]\n",
-      "File \u001b[0;32m~/cola/cola/algorithms/lanczos.py:186\u001b[0m, in \u001b[0;36mlanczos_parts\u001b[0;34m(A, rhs, max_iters, tol, pbar)\u001b[0m\n\u001b[1;32m    184\u001b[0m init_val \u001b[39m=\u001b[39m initialize_lanczos_vec(xnp, rhs, max_iters\u001b[39m=\u001b[39mmax_iters, dtype\u001b[39m=\u001b[39mA\u001b[39m.\u001b[39mdtype)\n\u001b[1;32m    185\u001b[0m while_fn, info \u001b[39m=\u001b[39m xnp\u001b[39m.\u001b[39mwhile_loop_winfo(error, pbar\u001b[39m=\u001b[39mpbar, tol\u001b[39m=\u001b[39mtol)\n\u001b[0;32m--> 186\u001b[0m state \u001b[39m=\u001b[39m while_fn(cond_fun, body_fun, init_val)\n\u001b[1;32m    187\u001b[0m i, vec, beta, alpha \u001b[39m=\u001b[39m state\n\u001b[1;32m    188\u001b[0m \u001b[39mreturn\u001b[39;00m alpha[\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m, \u001b[39m1\u001b[39m:], beta, vec[\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m, \u001b[39m1\u001b[39m:\u001b[39m-\u001b[39m\u001b[39m1\u001b[39m], i \u001b[39m-\u001b[39m \u001b[39m1\u001b[39m, info\n",
-      "File \u001b[0;32m~/cola/cola/utils/torch_tqdm.py:47\u001b[0m, in \u001b[0;36mwhile_loop_winfo.<locals>.new_while\u001b[0;34m(cond_fun, body_fun, init_val)\u001b[0m\n\u001b[1;32m     44\u001b[0m     info[\u001b[39m'\u001b[39m\u001b[39miterations\u001b[39m\u001b[39m'\u001b[39m] \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m \u001b[39m1\u001b[39m\n\u001b[1;32m     45\u001b[0m     \u001b[39mreturn\u001b[39;00m cond_fun(state)\n\u001b[0;32m---> 47\u001b[0m out \u001b[39m=\u001b[39m while_loop(newcond, body_fun, init_val)\n\u001b[1;32m     48\u001b[0m info[\u001b[39m'\u001b[39m\u001b[39miteration_time\u001b[39m\u001b[39m'\u001b[39m] \u001b[39m=\u001b[39m (time\u001b[39m.\u001b[39mtime() \u001b[39m-\u001b[39m info[\u001b[39m'\u001b[39m\u001b[39miteration_time\u001b[39m\u001b[39m'\u001b[39m]) \u001b[39m/\u001b[39m info[\u001b[39m'\u001b[39m\u001b[39miterations\u001b[39m\u001b[39m'\u001b[39m]\n\u001b[1;32m     49\u001b[0m error \u001b[39m=\u001b[39m errorfn(out)\n",
-      "File \u001b[0;32m~/cola/cola/utils/torch_tqdm.py:77\u001b[0m, in \u001b[0;36mwhile_loop\u001b[0;34m(cond_fun, body_fun, init_val)\u001b[0m\n\u001b[1;32m     75\u001b[0m val \u001b[39m=\u001b[39m init_val\n\u001b[1;32m     76\u001b[0m \u001b[39mwhile\u001b[39;00m cond_fun(val):\n\u001b[0;32m---> 77\u001b[0m     val \u001b[39m=\u001b[39m body_fun(val)\n\u001b[1;32m     78\u001b[0m \u001b[39mreturn\u001b[39;00m val\n",
-      "File \u001b[0;32m~/cola/cola/algorithms/lanczos.py:166\u001b[0m, in \u001b[0;36mlanczos_parts.<locals>.body_fun\u001b[0;34m(state)\u001b[0m\n\u001b[1;32m    164\u001b[0m aux \u001b[39m=\u001b[39m beta[\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m, [i \u001b[39m-\u001b[39m \u001b[39m1\u001b[39m]] \u001b[39m*\u001b[39m vec[\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m, i] \u001b[39m+\u001b[39m alpha[\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m, [i \u001b[39m-\u001b[39m \u001b[39m1\u001b[39m]] \u001b[39m*\u001b[39m vec[\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m, i \u001b[39m-\u001b[39m \u001b[39m1\u001b[39m]\n\u001b[1;32m    165\u001b[0m new_vec \u001b[39m-\u001b[39m\u001b[39m=\u001b[39m aux\n\u001b[0;32m--> 166\u001b[0m new_vec \u001b[39m=\u001b[39m do_double_gram(vec, new_vec, xnp)\n\u001b[1;32m    168\u001b[0m vec \u001b[39m=\u001b[39m xnp\u001b[39m.\u001b[39mupdate_array(vec, new_vec, \u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m, i \u001b[39m+\u001b[39m \u001b[39m1\u001b[39m)\n\u001b[1;32m    169\u001b[0m alpha \u001b[39m=\u001b[39m xnp\u001b[39m.\u001b[39mupdate_array(alpha, xnp\u001b[39m.\u001b[39mnorm(vec[\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m, i \u001b[39m+\u001b[39m \u001b[39m1\u001b[39m], axis\u001b[39m=\u001b[39m\u001b[39m-\u001b[39m\u001b[39m1\u001b[39m), \u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m, i)\n",
-      "File \u001b[0;32m~/cola/cola/algorithms/lanczos.py:202\u001b[0m, in \u001b[0;36mdo_double_gram\u001b[0;34m(vec, new_vec, xnp)\u001b[0m\n\u001b[1;32m    201\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mdo_double_gram\u001b[39m(vec, new_vec, xnp):\n\u001b[0;32m--> 202\u001b[0m     new_vec \u001b[39m=\u001b[39m do_gram(vec, new_vec, xnp)\n\u001b[1;32m    203\u001b[0m     new_vec \u001b[39m=\u001b[39m do_gram(vec, new_vec, xnp)\n\u001b[1;32m    204\u001b[0m     \u001b[39mreturn\u001b[39;00m new_vec\n",
-      "File \u001b[0;32m~/cola/cola/algorithms/lanczos.py:208\u001b[0m, in \u001b[0;36mdo_gram\u001b[0;34m(vec, new_vec, xnp)\u001b[0m\n\u001b[1;32m    207\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mdo_gram\u001b[39m(vec, new_vec, xnp):\n\u001b[0;32m--> 208\u001b[0m     aux \u001b[39m=\u001b[39m xnp\u001b[39m.\u001b[39;49msum(xnp\u001b[39m.\u001b[39;49mconj(vec) \u001b[39m*\u001b[39;49m xnp\u001b[39m.\u001b[39;49mexpand(new_vec, \u001b[39m-\u001b[39;49m\u001b[39m1\u001b[39;49m), axis\u001b[39m=\u001b[39;49m\u001b[39m-\u001b[39;49m\u001b[39m2\u001b[39;49m, keepdims\u001b[39m=\u001b[39;49m\u001b[39mTrue\u001b[39;49;00m)\n\u001b[1;32m    209\u001b[0m     new_vec \u001b[39m-\u001b[39m\u001b[39m=\u001b[39m xnp\u001b[39m.\u001b[39msum(vec \u001b[39m*\u001b[39m aux, axis\u001b[39m=\u001b[39m\u001b[39m-\u001b[39m\u001b[39m1\u001b[39m)\n\u001b[1;32m    210\u001b[0m     \u001b[39mreturn\u001b[39;00m new_vec\n",
-      "File \u001b[0;32m~/cola/cola/torch_fns.py:136\u001b[0m, in \u001b[0;36msum\u001b[0;34m(array, axis, keepdims)\u001b[0m\n\u001b[1;32m    132\u001b[0m     vec[loc] \u001b[39m=\u001b[39m \u001b[39m1.\u001b[39m\n\u001b[1;32m    133\u001b[0m     \u001b[39mreturn\u001b[39;00m vec\n\u001b[0;32m--> 136\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39msum\u001b[39m(array, axis\u001b[39m=\u001b[39m\u001b[39m0\u001b[39m, keepdims\u001b[39m=\u001b[39m\u001b[39mFalse\u001b[39;00m):\n\u001b[1;32m    137\u001b[0m     \u001b[39mreturn\u001b[39;00m torch\u001b[39m.\u001b[39msum(array, dim\u001b[39m=\u001b[39maxis, keepdims\u001b[39m=\u001b[39mkeepdims)\n\u001b[1;32m    140\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mpermute\u001b[39m(array, axes):\n",
-      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "tensor([-2.0163e-07,  7.4170e-07,  2.9827e-06]) tensor([-2.5722e-08,  5.0972e-07,  2.6101e-06])\n"
      ]
     }
    ],
    "source": [
-    "e,v = cola.eig(C,eig_slice=np.s_[:3])\n",
-    "print(e, torch.linalg.eigh(C.to_dense())[0][:3])\n",
-    "#TODO Andres: investigate why this is failing"
+    "e, v = cola.eig(C, eig_slice=np.s_[:3])\n",
+    "print(e, torch.linalg.eigh(C.to_dense())[0][:3])"
    ]
   }
  ],
diff --git a/docs/notebooks/colabs/Quick_Start.ipynb b/docs/notebooks/colabs/Quick_Start.ipynb
index dd8edadd..d5813bea 100644
--- a/docs/notebooks/colabs/Quick_Start.ipynb
+++ b/docs/notebooks/colabs/Quick_Start.ipynb
@@ -44,7 +44,7 @@
     {
      "data": {
       "text/plain": [
-       "<torch._C.Generator at 0x7fe235dde670>"
+       "<torch._C.Generator at 0x7f0f65a5a3f0>"
       ]
      },
      "execution_count": 1,
@@ -55,6 +55,10 @@
    "source": [
     "import cola\n",
     "import torch\n",
+    "import numpy as np\n",
+    "import warnings\n",
+    "\n",
+    "warnings.filterwarnings(\"ignore\")\n",
     "torch.manual_seed(21)"
    ]
   },
@@ -213,11 +217,10 @@
     }
    ],
    "source": [
-    "import numpy as np\n",
     "perm = torch.randperm(200)\n",
     "P = cola.ops.Permutation(perm)\n",
     "v = torch.randn(200)\n",
-    "print((P@v)[:5],v[perm[:5]])"
+    "print((P @ v)[:5], v[perm[:5]])"
    ]
   },
   {
@@ -263,8 +266,8 @@
     }
    ],
    "source": [
-    "D = cola.diag(torch.tensor([1,2,3,4.]))\n",
-    "DT = T[:4,:4]-D\n",
+    "D = cola.diag(torch.tensor([1, 2, 3, 4.]))\n",
+    "DT = T[:4, :4] - D\n",
     "print(DT.to_dense())"
    ]
   },
@@ -311,7 +314,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 10,
+   "execution_count": 9,
    "id": "06defdf2",
    "metadata": {},
    "outputs": [
@@ -329,13 +332,13 @@
     }
    ],
    "source": [
-    "K = cola.kron(A,D)\n",
+    "K = cola.kron(A, D)\n",
     "print(f\"{K[:5,:5].to_dense()}\\n shape: {K.shape}, object: {K}\")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 11,
+   "execution_count": 10,
    "id": "f660e95b",
    "metadata": {},
    "outputs": [
@@ -345,14 +348,14 @@
        "<2100x2100 BlockDiag[cola.ops.operators.Kronecker[cola.ops.operators.Dense, cola.ops.operators.Permutation], cola.ops.operators.Tridiagonal, cola.ops.operators.Diagonal] with dtype=torch.float32>"
       ]
      },
-     "execution_count": 11,
+     "execution_count": 10,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "d = 1/(1+torch.linspace(.2,10,1000))\n",
-    "X = cola.block_diag(cola.kron(A,P), T, cola.diag(d))\n",
+    "d = 1 / (1 + torch.linspace(.2, 10, 1000))\n",
+    "X = cola.block_diag(cola.kron(A, P), T, cola.diag(d))\n",
     "X"
    ]
   },
@@ -382,7 +385,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 24,
+   "execution_count": 11,
    "id": "4e803b0c",
    "metadata": {},
    "outputs": [
@@ -392,13 +395,13 @@
        "torch.Size([2100, 5])"
       ]
      },
-     "execution_count": 24,
+     "execution_count": 11,
      "metadata": {},
      "output_type": "execute_result"
     }
    ],
    "source": [
-    "X[:,:5].to_dense().shape"
+    "X[:, :5].to_dense().shape"
    ]
   },
   {
@@ -442,7 +445,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 21,
+   "execution_count": 14,
    "id": "04d88f0f",
    "metadata": {},
    "outputs": [
@@ -461,7 +464,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 44,
+   "execution_count": 15,
    "id": "b6432c5b",
    "metadata": {},
    "outputs": [
@@ -474,11 +477,11 @@
       "        [0.0932, 0.2152, 0.3083, 0.2139, 0.0864],\n",
       "        [0.0288, 0.0930, 0.2139, 0.3016, 0.1865],\n",
       "        [0.0066, 0.0275, 0.0864, 0.1865, 0.2153]])\n",
-      "scipy expm: [[0.30850828 0.21526888 0.09323523 0.02876081 0.0066488 ]\n",
-      " [0.21526891 0.30850452 0.2152388  0.09302244 0.02746145]\n",
-      " [0.09323531 0.2152388  0.30829173 0.21393944 0.08637362]\n",
-      " [0.02876082 0.09302243 0.2139395  0.30164295 0.18647799]\n",
-      " [0.0066488  0.02746145 0.08637366 0.18647805 0.21526927]]\n"
+      "scipy expm: [[0.30850834 0.21526882 0.09323522 0.02876082 0.0066488 ]\n",
+      " [0.21526891 0.30850452 0.21523876 0.09302247 0.02746145]\n",
+      " [0.09323531 0.21523881 0.30829167 0.21393949 0.08637363]\n",
+      " [0.02876082 0.09302244 0.21393946 0.30164295 0.18647799]\n",
+      " [0.0066488  0.02746145 0.08637364 0.18647805 0.21526927]]\n"
      ]
     }
    ],
@@ -490,42 +493,47 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 36,
+   "execution_count": 16,
    "id": "23e4504e",
    "metadata": {},
    "outputs": [
     {
-     "data": {
-      "text/plain": [
-       "tensor(0.+0.j)"
-      ]
-     },
-     "execution_count": 36,
-     "metadata": {},
-     "output_type": "execute_result"
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "CoLA exp(X): tensor([[0.3085, 0.2153, 0.0932, 0.0288, 0.0066],\n",
+      "        [0.2153, 0.3085, 0.2152, 0.0930, 0.0275],\n",
+      "        [0.0932, 0.2152, 0.3083, 0.2139, 0.0864],\n",
+      "        [0.0288, 0.0930, 0.2139, 0.3016, 0.1865],\n",
+      "        [0.0066, 0.0275, 0.0864, 0.1865, 0.2153]])\n",
+      "scipy expm: [[0.30850834 0.21526882 0.09323522 0.02876082 0.0066488 ]\n",
+      " [0.21526891 0.30850452 0.21523876 0.09302247 0.02746145]\n",
+      " [0.09323531 0.21523881 0.30829167 0.21393949 0.08637363]\n",
+      " [0.02876082 0.09302244 0.21393946 0.30164295 0.18647799]\n",
+      " [0.0066488  0.02746145 0.08637364 0.18647805 0.21526927]]\n"
+     ]
     }
    ],
    "source": [
-    "print(f\"CoLA exp(X): {cola.exp(T,method='iterative')[-5:,-5:].to_dense().real}\\nscipy expm: {scipy_expm[-5:,-5:]}\")\n",
-    "#TODO Andres: investigate why iterative exp not working for big matrices"
+    "print(f\"CoLA exp(X): {cola.exp(T, method='iterative')[-5:,-5:].to_dense().real}\\nscipy expm: {scipy_expm[-5:,-5:]}\")"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 54,
+   "execution_count": 17,
    "id": "32ad8a12",
    "metadata": {},
    "outputs": [],
    "source": [
-    "C = X.T@X\n",
+    "C = X.T @ X\n",
     "# let's annotate this matrix as positive definite to speed up the computation\n",
     "C = cola.PSD(C)\n",
-    "e,v = cola.eig(C,eig_slice=np.s_[-3:],max_iters=100)"
+    "e, v = cola.eig(C, eig_slice=np.s_[-3:], max_iters=100)"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": 55,
+   "execution_count": 18,
    "id": "2d0f3297",
    "metadata": {},
    "outputs": [
@@ -533,7 +541,7 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "tensor([15.9769, 15.9922, 15.9984]) tensor([15.9922, 15.9965, 15.9991])\n"
+      "tensor([15.9803, 15.9892, 15.9981]) tensor([15.9922, 15.9965, 15.9991])\n"
      ]
     }
    ],
@@ -543,47 +551,21 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 57,
+   "execution_count": 19,
    "id": "477eae7e",
    "metadata": {},
    "outputs": [
     {
-     "ename": "KeyboardInterrupt",
-     "evalue": "",
-     "output_type": "error",
-     "traceback": [
-      "\u001b[0;31m---------------------------------------------------------------------------\u001b[0m",
-      "\u001b[0;31mKeyboardInterrupt\u001b[0m                         Traceback (most recent call last)",
-      "Cell \u001b[0;32mIn[57], line 1\u001b[0m\n\u001b[0;32m----> 1\u001b[0m e,v \u001b[39m=\u001b[39m cola\u001b[39m.\u001b[39;49meig(C,eig_slice\u001b[39m=\u001b[39;49mnp\u001b[39m.\u001b[39;49ms_[:\u001b[39m3\u001b[39;49m])\n\u001b[1;32m      2\u001b[0m \u001b[39mprint\u001b[39m(e, torch\u001b[39m.\u001b[39mlinalg\u001b[39m.\u001b[39meigh(C\u001b[39m.\u001b[39mto_dense())[\u001b[39m0\u001b[39m][:\u001b[39m3\u001b[39m])\n",
-      "File \u001b[0;32m~/anaconda3/envs/cola/lib/python3.10/site-packages/plum/function.py:438\u001b[0m, in \u001b[0;36mFunction.__call__\u001b[0;34m(self, *args, **kw_args)\u001b[0m\n\u001b[1;32m    436\u001b[0m     method, return_type, loginfo \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mresolve_method(args, types)\n\u001b[1;32m    437\u001b[0m logging\u001b[39m.\u001b[39minfo(\u001b[39m\"\u001b[39m\u001b[39m%s\u001b[39;00m\u001b[39m\"\u001b[39m,loginfo)\n\u001b[0;32m--> 438\u001b[0m \u001b[39mreturn\u001b[39;00m _convert(method(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkw_args), return_type)\n",
-      "File \u001b[0;32m~/cola/cola/linalg/eigs.py:74\u001b[0m, in \u001b[0;36meig\u001b[0;34m(A, **kwargs)\u001b[0m\n\u001b[1;32m     72\u001b[0m     \u001b[39mreturn\u001b[39;00m eig_vals[eig_slice], Unitary(lazify(eig_vecs[:, eig_slice]))\n\u001b[1;32m     73\u001b[0m \u001b[39melif\u001b[39;00m method \u001b[39min\u001b[39;00m (\u001b[39m'\u001b[39m\u001b[39mlanczos\u001b[39m\u001b[39m'\u001b[39m, \u001b[39m'\u001b[39m\u001b[39miterative\u001b[39m\u001b[39m'\u001b[39m) \u001b[39mor\u001b[39;00m (method \u001b[39m==\u001b[39m \u001b[39m'\u001b[39m\u001b[39mauto\u001b[39m\u001b[39m'\u001b[39m \u001b[39mand\u001b[39;00m prod(A\u001b[39m.\u001b[39mshape) \u001b[39m>\u001b[39m\u001b[39m=\u001b[39m \u001b[39m1e6\u001b[39m):\n\u001b[0;32m---> 74\u001b[0m     eig_vals, eig_vecs \u001b[39m=\u001b[39m eig(LanczosDecomposition(A, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkws), eig_slice\u001b[39m=\u001b[39;49meig_slice)\n\u001b[1;32m     75\u001b[0m     \u001b[39mreturn\u001b[39;00m eig_vals, eig_vecs\n\u001b[1;32m     76\u001b[0m \u001b[39melse\u001b[39;00m:\n",
-      "File \u001b[0;32m~/anaconda3/envs/cola/lib/python3.10/site-packages/plum/function.py:438\u001b[0m, in \u001b[0;36mFunction.__call__\u001b[0;34m(self, *args, **kw_args)\u001b[0m\n\u001b[1;32m    436\u001b[0m     method, return_type, loginfo \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mresolve_method(args, types)\n\u001b[1;32m    437\u001b[0m logging\u001b[39m.\u001b[39minfo(\u001b[39m\"\u001b[39m\u001b[39m%s\u001b[39;00m\u001b[39m\"\u001b[39m,loginfo)\n\u001b[0;32m--> 438\u001b[0m \u001b[39mreturn\u001b[39;00m _convert(method(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkw_args), return_type)\n",
-      "File \u001b[0;32m~/cola/cola/decompositions.py:51\u001b[0m, in \u001b[0;36meig\u001b[0;34m(QH, **kwargs)\u001b[0m\n\u001b[1;32m     48\u001b[0m \u001b[39m@eig\u001b[39m\u001b[39m.\u001b[39mdispatch\n\u001b[1;32m     49\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39meig\u001b[39m(QH: UnitaryDecomposition, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs):\n\u001b[1;32m     50\u001b[0m     Q, H, xnp \u001b[39m=\u001b[39m QH\u001b[39m.\u001b[39mQ, QH\u001b[39m.\u001b[39mM, QH\u001b[39m.\u001b[39mxnp\n\u001b[0;32m---> 51\u001b[0m     eig_vals, eig_vecs \u001b[39m=\u001b[39m eig(H, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m     52\u001b[0m     eig_vecs \u001b[39m=\u001b[39m xnp\u001b[39m.\u001b[39mcast(Q\u001b[39m.\u001b[39mto_dense(), dtype\u001b[39m=\u001b[39meig_vecs\u001b[39m.\u001b[39mdtype) \u001b[39m@\u001b[39m eig_vecs\u001b[39m.\u001b[39mto_dense()\n\u001b[1;32m     53\u001b[0m     eig_vecs \u001b[39m=\u001b[39m Unitary(lazify(eig_vecs))\n",
-      "    \u001b[0;31m[... skipping similar frames: Function.__call__ at line 438 (1 times)]\u001b[0m\n",
-      "File \u001b[0;32m~/cola/cola/linalg/eigs.py:74\u001b[0m, in \u001b[0;36meig\u001b[0;34m(A, **kwargs)\u001b[0m\n\u001b[1;32m     72\u001b[0m     \u001b[39mreturn\u001b[39;00m eig_vals[eig_slice], Unitary(lazify(eig_vecs[:, eig_slice]))\n\u001b[1;32m     73\u001b[0m \u001b[39melif\u001b[39;00m method \u001b[39min\u001b[39;00m (\u001b[39m'\u001b[39m\u001b[39mlanczos\u001b[39m\u001b[39m'\u001b[39m, \u001b[39m'\u001b[39m\u001b[39miterative\u001b[39m\u001b[39m'\u001b[39m) \u001b[39mor\u001b[39;00m (method \u001b[39m==\u001b[39m \u001b[39m'\u001b[39m\u001b[39mauto\u001b[39m\u001b[39m'\u001b[39m \u001b[39mand\u001b[39;00m prod(A\u001b[39m.\u001b[39mshape) \u001b[39m>\u001b[39m\u001b[39m=\u001b[39m \u001b[39m1e6\u001b[39m):\n\u001b[0;32m---> 74\u001b[0m     eig_vals, eig_vecs \u001b[39m=\u001b[39m eig(LanczosDecomposition(A, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkws), eig_slice\u001b[39m=\u001b[39;49meig_slice)\n\u001b[1;32m     75\u001b[0m     \u001b[39mreturn\u001b[39;00m eig_vals, eig_vecs\n\u001b[1;32m     76\u001b[0m \u001b[39melse\u001b[39;00m:\n",
-      "    \u001b[0;31m[... skipping similar frames: Function.__call__ at line 438 (1 times)]\u001b[0m\n",
-      "File \u001b[0;32m~/cola/cola/decompositions.py:51\u001b[0m, in \u001b[0;36meig\u001b[0;34m(QH, **kwargs)\u001b[0m\n\u001b[1;32m     48\u001b[0m \u001b[39m@eig\u001b[39m\u001b[39m.\u001b[39mdispatch\n\u001b[1;32m     49\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39meig\u001b[39m(QH: UnitaryDecomposition, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs):\n\u001b[1;32m     50\u001b[0m     Q, H, xnp \u001b[39m=\u001b[39m QH\u001b[39m.\u001b[39mQ, QH\u001b[39m.\u001b[39mM, QH\u001b[39m.\u001b[39mxnp\n\u001b[0;32m---> 51\u001b[0m     eig_vals, eig_vecs \u001b[39m=\u001b[39m eig(H, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m     52\u001b[0m     eig_vecs \u001b[39m=\u001b[39m xnp\u001b[39m.\u001b[39mcast(Q\u001b[39m.\u001b[39mto_dense(), dtype\u001b[39m=\u001b[39meig_vecs\u001b[39m.\u001b[39mdtype) \u001b[39m@\u001b[39m eig_vecs\u001b[39m.\u001b[39mto_dense()\n\u001b[1;32m     53\u001b[0m     eig_vecs \u001b[39m=\u001b[39m Unitary(lazify(eig_vecs))\n",
-      "    \u001b[0;31m[... skipping similar frames: Function.__call__ at line 438 (156 times), eig at line 74 (78 times), eig at line 51 (77 times)]\u001b[0m\n",
-      "File \u001b[0;32m~/cola/cola/decompositions.py:51\u001b[0m, in \u001b[0;36meig\u001b[0;34m(QH, **kwargs)\u001b[0m\n\u001b[1;32m     48\u001b[0m \u001b[39m@eig\u001b[39m\u001b[39m.\u001b[39mdispatch\n\u001b[1;32m     49\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39meig\u001b[39m(QH: UnitaryDecomposition, \u001b[39m*\u001b[39m\u001b[39m*\u001b[39mkwargs):\n\u001b[1;32m     50\u001b[0m     Q, H, xnp \u001b[39m=\u001b[39m QH\u001b[39m.\u001b[39mQ, QH\u001b[39m.\u001b[39mM, QH\u001b[39m.\u001b[39mxnp\n\u001b[0;32m---> 51\u001b[0m     eig_vals, eig_vecs \u001b[39m=\u001b[39m eig(H, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkwargs)\n\u001b[1;32m     52\u001b[0m     eig_vecs \u001b[39m=\u001b[39m xnp\u001b[39m.\u001b[39mcast(Q\u001b[39m.\u001b[39mto_dense(), dtype\u001b[39m=\u001b[39meig_vecs\u001b[39m.\u001b[39mdtype) \u001b[39m@\u001b[39m eig_vecs\u001b[39m.\u001b[39mto_dense()\n\u001b[1;32m     53\u001b[0m     eig_vecs \u001b[39m=\u001b[39m Unitary(lazify(eig_vecs))\n",
-      "File \u001b[0;32m~/anaconda3/envs/cola/lib/python3.10/site-packages/plum/function.py:438\u001b[0m, in \u001b[0;36mFunction.__call__\u001b[0;34m(self, *args, **kw_args)\u001b[0m\n\u001b[1;32m    436\u001b[0m     method, return_type, loginfo \u001b[39m=\u001b[39m \u001b[39mself\u001b[39m\u001b[39m.\u001b[39mresolve_method(args, types)\n\u001b[1;32m    437\u001b[0m logging\u001b[39m.\u001b[39minfo(\u001b[39m\"\u001b[39m\u001b[39m%s\u001b[39;00m\u001b[39m\"\u001b[39m,loginfo)\n\u001b[0;32m--> 438\u001b[0m \u001b[39mreturn\u001b[39;00m _convert(method(\u001b[39m*\u001b[39;49margs, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkw_args), return_type)\n",
-      "File \u001b[0;32m~/cola/cola/linalg/eigs.py:74\u001b[0m, in \u001b[0;36meig\u001b[0;34m(A, **kwargs)\u001b[0m\n\u001b[1;32m     72\u001b[0m     \u001b[39mreturn\u001b[39;00m eig_vals[eig_slice], Unitary(lazify(eig_vecs[:, eig_slice]))\n\u001b[1;32m     73\u001b[0m \u001b[39melif\u001b[39;00m method \u001b[39min\u001b[39;00m (\u001b[39m'\u001b[39m\u001b[39mlanczos\u001b[39m\u001b[39m'\u001b[39m, \u001b[39m'\u001b[39m\u001b[39miterative\u001b[39m\u001b[39m'\u001b[39m) \u001b[39mor\u001b[39;00m (method \u001b[39m==\u001b[39m \u001b[39m'\u001b[39m\u001b[39mauto\u001b[39m\u001b[39m'\u001b[39m \u001b[39mand\u001b[39;00m prod(A\u001b[39m.\u001b[39mshape) \u001b[39m>\u001b[39m\u001b[39m=\u001b[39m \u001b[39m1e6\u001b[39m):\n\u001b[0;32m---> 74\u001b[0m     eig_vals, eig_vecs \u001b[39m=\u001b[39m eig(LanczosDecomposition(A, \u001b[39m*\u001b[39;49m\u001b[39m*\u001b[39;49mkws), eig_slice\u001b[39m=\u001b[39meig_slice)\n\u001b[1;32m     75\u001b[0m     \u001b[39mreturn\u001b[39;00m eig_vals, eig_vecs\n\u001b[1;32m     76\u001b[0m \u001b[39melse\u001b[39;00m:\n",
-      "File \u001b[0;32m~/cola/cola/algorithms/lanczos.py:103\u001b[0m, in \u001b[0;36mLanczosDecomposition\u001b[0;34m(A, start_vector, max_iters, tol, pbar)\u001b[0m\n\u001b[1;32m    100\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mLanczosDecomposition\u001b[39m(A: LinearOperator, start_vector\u001b[39m=\u001b[39m\u001b[39mNone\u001b[39;00m, max_iters\u001b[39m=\u001b[39m\u001b[39m100\u001b[39m, tol\u001b[39m=\u001b[39m\u001b[39m1e-7\u001b[39m, pbar\u001b[39m=\u001b[39m\u001b[39mFalse\u001b[39;00m):\n\u001b[1;32m    101\u001b[0m \u001b[39m    \u001b[39m\u001b[39m\"\"\" Provides the Lanczos decomposition of a matrix A = Q T Q^*.\u001b[39;00m\n\u001b[1;32m    102\u001b[0m \u001b[39m    LinearOperator form of lanczos, see lanczos for arguments.\"\"\"\u001b[39;00m\n\u001b[0;32m--> 103\u001b[0m     Q, T, info \u001b[39m=\u001b[39m lanczos_decomp(A\u001b[39m=\u001b[39;49mA, start_vector\u001b[39m=\u001b[39;49mstart_vector, max_iters\u001b[39m=\u001b[39;49mmax_iters, tol\u001b[39m=\u001b[39;49mtol,\n\u001b[1;32m    104\u001b[0m                                 pbar\u001b[39m=\u001b[39;49mpbar)\n\u001b[1;32m    105\u001b[0m     A_approx \u001b[39m=\u001b[39m cola\u001b[39m.\u001b[39mUnitaryDecomposition(lazify(Q), SelfAdjoint(lazify(T)))\n\u001b[1;32m    106\u001b[0m     A_approx\u001b[39m.\u001b[39minfo \u001b[39m=\u001b[39m info\n",
-      "File \u001b[0;32m~/cola/cola/algorithms/lanczos.py:132\u001b[0m, in \u001b[0;36mlanczos_decomp\u001b[0;34m(A, start_vector, max_iters, tol, pbar)\u001b[0m\n\u001b[1;32m    130\u001b[0m \u001b[39mif\u001b[39;00m start_vector \u001b[39mis\u001b[39;00m \u001b[39mNone\u001b[39;00m:\n\u001b[1;32m    131\u001b[0m     start_vector \u001b[39m=\u001b[39m xnp\u001b[39m.\u001b[39mfixed_normal_samples((A\u001b[39m.\u001b[39mshape[\u001b[39m0\u001b[39m], \u001b[39m1\u001b[39m))\n\u001b[0;32m--> 132\u001b[0m alpha, beta, vec, iters, info \u001b[39m=\u001b[39m lanczos_parts(A\u001b[39m=\u001b[39;49mA, rhs\u001b[39m=\u001b[39;49mstart_vector, max_iters\u001b[39m=\u001b[39;49mmax_iters,\n\u001b[1;32m    133\u001b[0m                                               tol\u001b[39m=\u001b[39;49mtol, pbar\u001b[39m=\u001b[39;49mpbar)\n\u001b[1;32m    134\u001b[0m alpha, beta, Q \u001b[39m=\u001b[39m alpha[\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m, :iters \u001b[39m-\u001b[39m \u001b[39m1\u001b[39m], beta[\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m, :iters], vec[\u001b[39m0\u001b[39m, :, :iters]\n\u001b[1;32m    135\u001b[0m T \u001b[39m=\u001b[39m construct_tridiagonal_batched(alpha, beta, alpha)[\u001b[39m0\u001b[39m]\n",
-      "File \u001b[0;32m~/cola/cola/algorithms/lanczos.py:186\u001b[0m, in \u001b[0;36mlanczos_parts\u001b[0;34m(A, rhs, max_iters, tol, pbar)\u001b[0m\n\u001b[1;32m    184\u001b[0m init_val \u001b[39m=\u001b[39m initialize_lanczos_vec(xnp, rhs, max_iters\u001b[39m=\u001b[39mmax_iters, dtype\u001b[39m=\u001b[39mA\u001b[39m.\u001b[39mdtype)\n\u001b[1;32m    185\u001b[0m while_fn, info \u001b[39m=\u001b[39m xnp\u001b[39m.\u001b[39mwhile_loop_winfo(error, pbar\u001b[39m=\u001b[39mpbar, tol\u001b[39m=\u001b[39mtol)\n\u001b[0;32m--> 186\u001b[0m state \u001b[39m=\u001b[39m while_fn(cond_fun, body_fun, init_val)\n\u001b[1;32m    187\u001b[0m i, vec, beta, alpha \u001b[39m=\u001b[39m state\n\u001b[1;32m    188\u001b[0m \u001b[39mreturn\u001b[39;00m alpha[\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m, \u001b[39m1\u001b[39m:], beta, vec[\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m, \u001b[39m1\u001b[39m:\u001b[39m-\u001b[39m\u001b[39m1\u001b[39m], i \u001b[39m-\u001b[39m \u001b[39m1\u001b[39m, info\n",
-      "File \u001b[0;32m~/cola/cola/utils/torch_tqdm.py:47\u001b[0m, in \u001b[0;36mwhile_loop_winfo.<locals>.new_while\u001b[0;34m(cond_fun, body_fun, init_val)\u001b[0m\n\u001b[1;32m     44\u001b[0m     info[\u001b[39m'\u001b[39m\u001b[39miterations\u001b[39m\u001b[39m'\u001b[39m] \u001b[39m+\u001b[39m\u001b[39m=\u001b[39m \u001b[39m1\u001b[39m\n\u001b[1;32m     45\u001b[0m     \u001b[39mreturn\u001b[39;00m cond_fun(state)\n\u001b[0;32m---> 47\u001b[0m out \u001b[39m=\u001b[39m while_loop(newcond, body_fun, init_val)\n\u001b[1;32m     48\u001b[0m info[\u001b[39m'\u001b[39m\u001b[39miteration_time\u001b[39m\u001b[39m'\u001b[39m] \u001b[39m=\u001b[39m (time\u001b[39m.\u001b[39mtime() \u001b[39m-\u001b[39m info[\u001b[39m'\u001b[39m\u001b[39miteration_time\u001b[39m\u001b[39m'\u001b[39m]) \u001b[39m/\u001b[39m info[\u001b[39m'\u001b[39m\u001b[39miterations\u001b[39m\u001b[39m'\u001b[39m]\n\u001b[1;32m     49\u001b[0m error \u001b[39m=\u001b[39m errorfn(out)\n",
-      "File \u001b[0;32m~/cola/cola/utils/torch_tqdm.py:77\u001b[0m, in \u001b[0;36mwhile_loop\u001b[0;34m(cond_fun, body_fun, init_val)\u001b[0m\n\u001b[1;32m     75\u001b[0m val \u001b[39m=\u001b[39m init_val\n\u001b[1;32m     76\u001b[0m \u001b[39mwhile\u001b[39;00m cond_fun(val):\n\u001b[0;32m---> 77\u001b[0m     val \u001b[39m=\u001b[39m body_fun(val)\n\u001b[1;32m     78\u001b[0m \u001b[39mreturn\u001b[39;00m val\n",
-      "File \u001b[0;32m~/cola/cola/algorithms/lanczos.py:166\u001b[0m, in \u001b[0;36mlanczos_parts.<locals>.body_fun\u001b[0;34m(state)\u001b[0m\n\u001b[1;32m    164\u001b[0m aux \u001b[39m=\u001b[39m beta[\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m, [i \u001b[39m-\u001b[39m \u001b[39m1\u001b[39m]] \u001b[39m*\u001b[39m vec[\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m, i] \u001b[39m+\u001b[39m alpha[\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m, [i \u001b[39m-\u001b[39m \u001b[39m1\u001b[39m]] \u001b[39m*\u001b[39m vec[\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m, i \u001b[39m-\u001b[39m \u001b[39m1\u001b[39m]\n\u001b[1;32m    165\u001b[0m new_vec \u001b[39m-\u001b[39m\u001b[39m=\u001b[39m aux\n\u001b[0;32m--> 166\u001b[0m new_vec \u001b[39m=\u001b[39m do_double_gram(vec, new_vec, xnp)\n\u001b[1;32m    168\u001b[0m vec \u001b[39m=\u001b[39m xnp\u001b[39m.\u001b[39mupdate_array(vec, new_vec, \u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m, i \u001b[39m+\u001b[39m \u001b[39m1\u001b[39m)\n\u001b[1;32m    169\u001b[0m alpha \u001b[39m=\u001b[39m xnp\u001b[39m.\u001b[39mupdate_array(alpha, xnp\u001b[39m.\u001b[39mnorm(vec[\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m, i \u001b[39m+\u001b[39m \u001b[39m1\u001b[39m], axis\u001b[39m=\u001b[39m\u001b[39m-\u001b[39m\u001b[39m1\u001b[39m), \u001b[39m.\u001b[39m\u001b[39m.\u001b[39m\u001b[39m.\u001b[39m, i)\n",
-      "File \u001b[0;32m~/cola/cola/algorithms/lanczos.py:202\u001b[0m, in \u001b[0;36mdo_double_gram\u001b[0;34m(vec, new_vec, xnp)\u001b[0m\n\u001b[1;32m    201\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mdo_double_gram\u001b[39m(vec, new_vec, xnp):\n\u001b[0;32m--> 202\u001b[0m     new_vec \u001b[39m=\u001b[39m do_gram(vec, new_vec, xnp)\n\u001b[1;32m    203\u001b[0m     new_vec \u001b[39m=\u001b[39m do_gram(vec, new_vec, xnp)\n\u001b[1;32m    204\u001b[0m     \u001b[39mreturn\u001b[39;00m new_vec\n",
-      "File \u001b[0;32m~/cola/cola/algorithms/lanczos.py:208\u001b[0m, in \u001b[0;36mdo_gram\u001b[0;34m(vec, new_vec, xnp)\u001b[0m\n\u001b[1;32m    207\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mdo_gram\u001b[39m(vec, new_vec, xnp):\n\u001b[0;32m--> 208\u001b[0m     aux \u001b[39m=\u001b[39m xnp\u001b[39m.\u001b[39;49msum(xnp\u001b[39m.\u001b[39;49mconj(vec) \u001b[39m*\u001b[39;49m xnp\u001b[39m.\u001b[39;49mexpand(new_vec, \u001b[39m-\u001b[39;49m\u001b[39m1\u001b[39;49m), axis\u001b[39m=\u001b[39;49m\u001b[39m-\u001b[39;49m\u001b[39m2\u001b[39;49m, keepdims\u001b[39m=\u001b[39;49m\u001b[39mTrue\u001b[39;49;00m)\n\u001b[1;32m    209\u001b[0m     new_vec \u001b[39m-\u001b[39m\u001b[39m=\u001b[39m xnp\u001b[39m.\u001b[39msum(vec \u001b[39m*\u001b[39m aux, axis\u001b[39m=\u001b[39m\u001b[39m-\u001b[39m\u001b[39m1\u001b[39m)\n\u001b[1;32m    210\u001b[0m     \u001b[39mreturn\u001b[39;00m new_vec\n",
-      "File \u001b[0;32m~/cola/cola/torch_fns.py:136\u001b[0m, in \u001b[0;36msum\u001b[0;34m(array, axis, keepdims)\u001b[0m\n\u001b[1;32m    132\u001b[0m     vec[loc] \u001b[39m=\u001b[39m \u001b[39m1.\u001b[39m\n\u001b[1;32m    133\u001b[0m     \u001b[39mreturn\u001b[39;00m vec\n\u001b[0;32m--> 136\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39msum\u001b[39m(array, axis\u001b[39m=\u001b[39m\u001b[39m0\u001b[39m, keepdims\u001b[39m=\u001b[39m\u001b[39mFalse\u001b[39;00m):\n\u001b[1;32m    137\u001b[0m     \u001b[39mreturn\u001b[39;00m torch\u001b[39m.\u001b[39msum(array, dim\u001b[39m=\u001b[39maxis, keepdims\u001b[39m=\u001b[39mkeepdims)\n\u001b[1;32m    140\u001b[0m \u001b[39mdef\u001b[39;00m \u001b[39mpermute\u001b[39m(array, axes):\n",
-      "\u001b[0;31mKeyboardInterrupt\u001b[0m: "
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "tensor([-2.0163e-07,  7.4170e-07,  2.9827e-06]) tensor([-2.5722e-08,  5.0972e-07,  2.6101e-06])\n"
      ]
     }
    ],
    "source": [
-    "e,v = cola.eig(C,eig_slice=np.s_[:3])\n",
-    "print(e, torch.linalg.eigh(C.to_dense())[0][:3])\n",
-    "#TODO Andres: investigate why this is failing"
+    "e, v = cola.eig(C, eig_slice=np.s_[:3])\n",
+    "print(e, torch.linalg.eigh(C.to_dense())[0][:3])"
    ]
   }
  ],