fix: remove redundant reshape (#275)

amazon-braket · Dec 6, 2024 · f1f9355 · f1f9355
1 parent cba0608
commit f1f9355
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 4 deletions.
diff --git a/src/braket/default_simulator/density_matrix_simulation.py b/src/braket/default_simulator/density_matrix_simulation.py
@@ -185,11 +185,11 @@ def _apply_gate(
             np.ndarray: output density matrix
         """
         # left product
-        state = multiply_matrix(state, np.reshape(matrix, [2] * len(targets) * 2), targets)
+        state = multiply_matrix(state, matrix, targets)
         # right product
         state = multiply_matrix(
             state,
-            np.reshape(matrix.conjugate(), [2] * len(targets) * 2),
+            matrix.conjugate(),
             tuple(i + qubit_count for i in targets),
         )
         return state

diff --git a/src/braket/default_simulator/linalg_utils.py b/src/braket/default_simulator/linalg_utils.py
@@ -83,9 +83,8 @@ def _multiply_matrix(
 
     # Axes given in `operation.targets` are in the first positions.
     unused_idxs = [idx for idx in range(len(state.shape)) if idx not in targets]
-    permutation = list(targets) + unused_idxs
     # Invert the permutation to put the indices in the correct place
-    inverse_permutation = np.argsort(permutation)
+    inverse_permutation = np.argsort([*targets, *unused_idxs])
     return np.transpose(product, inverse_permutation)