Fixing Grad calc

BQSKit · Sep 27, 2024 · 6e02414 · 6e02414
1 parent 779c8ab
commit 6e02414
Showing 1 changed file with 9 additions and 8 deletions.
diff --git a/bqskit/ir/gates/parameterized/diagonal.py b/bqskit/ir/gates/parameterized/diagonal.py
@@ -51,16 +51,17 @@ def get_grad(self, params: RealVector = []) -> npt.NDArray[np.complex128]:
         """
         self.check_parameters(params)
 
-        mat = np.eye(2 ** self.num_qudits, dtype=np.complex128)
+        grad = np.zeros(
+            (
+                len(params), 2 ** self.num_qudits,
+                2 ** self.num_qudits,
+            ), dtype=np.complex128,
+        )
 
-        for i in range(1, 2 ** self.num_qudits):
-            mat[i][i] = 1j * np.exp(1j * params[i - 1])
+        for i, ind in enumerate(range(1, 2 ** self.num_qudits)):
+            grad[ind][i][i] = 1j * np.exp(1j * params[ind])
 
-        return np.array(
-            [
-                mat,
-            ], dtype=np.complex128,
-        )
+        return grad
 
     def optimize(self, env_matrix: npt.NDArray[np.complex128]) -> list[float]:
         """