fixed einsum string (#592)

Author: ziofil

mrmustard/physics/ansatz/polyexp_ansatz.py

@@ -223,8 +223,8 @@ def PS(self) -> PolyExpAnsatz:
         if self.num_derived_vars == 0:
             W = math.conj(math.rotmat(n // 2)) / math.sqrt(settings.HBAR, dtype=math.complex128)
 
-            A = math.einsum("ji, ...jk, kl-> ...il", W, self.A, W)
-            b = math.einsum("ij, ...j-> ...i", W, self.b)
+            A = math.einsum("ji,...jk,kl->...il", W, self.A, W)
+            b = math.einsum("ij,...j->...i", W, self.b)
             c = self.c / (2 * settings.HBAR) ** (n // 2)
             return PolyExpAnsatz(A, b, c, lin_sup=self._lin_sup)
 
@@ -241,8 +241,8 @@ def PS(self) -> PolyExpAnsatz:
 
             W = math.conj(math.rotmat(n // 2)) / math.sqrt(settings.HBAR, dtype=math.complex128)
 
-            A = math.einsum("ji, ...jk, kl-> ...il", W, A_tmp, W)
-            b = math.einsum("ij, ...j-> ...i", W, b)
+            A = math.einsum("ji,...jk,kl->...il", W, A_tmp, W)
+            b = math.einsum("ij,...j->...i", W, b)
             c = c / (2 * settings.HBAR)
 
             A_final = A[..., [0, 2, 1, 3], :][..., :, [0, 2, 1, 3]]