diff --git a/docs/source/quantum_five_qubit.ipynb b/docs/source/quantum_five_qubit.ipynb
index 784492e..c558824 100644
--- a/docs/source/quantum_five_qubit.ipynb
+++ b/docs/source/quantum_five_qubit.ipynb
@@ -64,15 +64,15 @@
    "outputs": [],
    "source": [
     "stabilizer_generators = [\n",
-    "    \"XZZXI\", # 1001011000 -> 0, 3, 5, 6\n",
-    "    \"IXZZX\", # 0010010110 -> 2, 5, 6, 7\n",
-    "    \"XIXZZ\", # 1000100101 -> 0, 4, 7, 8\n",
-    "    \"ZXIXZ\", # 0110001001 -> 1, 2, 6, 8\n",
+    "    \"XZZXI\",  # 1001011000 -> 0, 3, 5, 6 -> 2, 5, 7, 8\n",
+    "    \"IXZZX\",  # 0010010110 -> 2, 5, 7, 8 -> 4, 7, 9, 10\n",
+    "    \"XIXZZ\",  # 1000100101 -> 0, 4, 7, 9 -> 2, 6, 9, 11\n",
+    "    \"ZXIXZ\",  # 0110001001 -> 1, 2, 6, 9 -> 3, 4, 8, 11\n",
     "]\n",
     "\n",
     "logical_operators = [\n",
-    "    \"XXXXX\", # X -> 0101010101\n",
-    "    \"ZZZZZ\", # Z -> 1010101010\n",
+    "    \"XXXXX\",  # X -> 0101010101 -> 1, 3, 5, 7, 9\n",
+    "    \"ZZZZZ\",  # Z -> 1010101010 -> 0, 2, 4, 6, 8\n",
     "]"
    ]
   },
@@ -84,7 +84,7 @@
     {
      "data": {
       "text/plain": [
-       "[[2, 4, 6, 8], [4, 6, 8, 10], [2, 6, 8, 10], [2, 4, 8, 10]]"
+       "[[2, 5, 7, 8], [4, 7, 9, 10], [2, 6, 9, 11], [3, 4, 8, 11]]"
       ]
      },
      "execution_count": 3,
@@ -98,9 +98,23 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
    "metadata": {},
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "text/plain": [
+       "[[[2], [5, 7], [3, 4, 6], [8]],\n",
+       " [[4], [7, 9], [5, 6, 8], [10]],\n",
+       " [[2], [6, 9], [3, 4, 5, 7, 8, 10], [11]],\n",
+       " [[3], [4, 8], [5, 6, 7, 9, 10], [11]]]"
+      ]
+     },
+     "execution_count": 4,
+     "metadata": {},
+     "output_type": "execute_result"
+    }
+   ],
    "source": [
     "custom_code_constraint_sites(stabilizer_generators, logical_operators)"
    ]
@@ -111,7 +125,7 @@
    "metadata": {},
    "outputs": [],
    "source": [
-    "custom_code_logicals(stabilizer_generators, logical_operators)"
+    "custom_code_logicals(logical_operators)"
    ]
   },
   {
diff --git a/docs/source/shor.ipynb b/docs/source/shor.ipynb
index 8608148..cc7cc70 100644
--- a/docs/source/shor.ipynb
+++ b/docs/source/shor.ipynb
@@ -351,10 +351,10 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "100%|██████████| 2/2 [00:00<00:00, 196.19it/s]\n",
-      "100%|██████████| 6/6 [00:00<00:00, 547.94it/s]\n",
-      "100%|██████████| 1/1 [00:00<00:00, 181.89it/s]\n",
-      "100%|██████████| 1/1 [00:00<00:00, 132.59it/s]\n"
+      "100%|██████████| 2/2 [00:00<00:00, 194.04it/s]\n",
+      "100%|██████████| 6/6 [00:00<00:00, 506.47it/s]\n",
+      "100%|██████████| 1/1 [00:00<00:00, 157.70it/s]\n",
+      "100%|██████████| 1/1 [00:00<00:00, 123.62it/s]\n"
      ]
     }
    ],
@@ -454,13 +454,20 @@
      "text": [
       "0it [00:00, ?it/s]/Users/aleksandrberezutskii/mdopt/mdopt/mps/canonical.py:274: RuntimeWarning: invalid value encountered in divide\n",
       "  dense /= np.linalg.norm(dense, ord=norm)\n",
-      "27it [00:00, 61.83it/s]\n"
+      "27it [00:00, 57.28it/s]\n"
      ]
     }
    ],
    "source": [
     "one_qubit_outputs = [\n",
-    "    decode_css(code, error, bias_type=\"Bitflip\", renormalise=renormalise, silent=True, bias_prob=0)\n",
+    "    decode_css(\n",
+    "        code,\n",
+    "        error,\n",
+    "        bias_type=\"Bitflip\",\n",
+    "        renormalise=renormalise,\n",
+    "        silent=True,\n",
+    "        bias_prob=0,\n",
+    "    )\n",
     "    for error in tqdm(one_qubit_paulis)\n",
     "]\n",
     "one_qubit_corrections_distribution = [output[0] for output in one_qubit_outputs]"
@@ -475,7 +482,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "324it [00:04, 74.09it/s]\n"
+      "324it [00:04, 71.60it/s]\n"
      ]
     }
    ],
@@ -496,7 +503,7 @@
      "name": "stderr",
      "output_type": "stream",
      "text": [
-      "2268it [00:31, 70.93it/s]\n"
+      "2268it [00:32, 68.75it/s]\n"
      ]
     }
    ],