fixing tests

PennyLaneAI · Oct 17, 2023 · a7036a9 · a7036a9
1 parent 906b539
commit a7036a9
Show file tree

Hide file tree

Showing 2 changed files with 5 additions and 2 deletions.
diff --git a/pennylane/templates/state_preparations/cosine_window.py b/pennylane/templates/state_preparations/cosine_window.py
@@ -95,7 +95,7 @@ def compute_decomposition(wires):
 
     def state_vector(self, wire_order=None):
         num_op_wires = len(self.wires)
-        op_vector_shape = (-1,) + (2,) * num_op_wires
+        op_vector_shape = (-1,) + (2,) * num_op_wires if self.batch_size else (2,) * num_op_wires
         vector = np.array(
             [
                 np.sqrt(2)
@@ -118,6 +118,9 @@ def state_vector(self, wire_order=None):
             [Ellipsis] + [slice(None)] * num_op_wires + [0] * (num_total_wires - num_op_wires)
         )
         ket_shape = [2] * num_total_wires
+        if self.batch_size:
+            # Add broadcasted dimension to the shape of the state vector
+            ket_shape = [self.batch_size] + ket_shape
 
         ket = np.zeros(ket_shape, dtype=np.complex128)
         ket[indices] = op_vector
@@ -130,7 +133,6 @@ def state_vector(self, wire_order=None):
                 # If the operation is broadcasted, the desired order must include the batch dimension
                 # as the first dimension.
                 desired_order = [0] + [d + 1 for d in desired_order]
-
             ket = ket.transpose(desired_order)
 
         return math.convert_like(ket, op_vector)
diff --git a/tests/templates/test_state_preparations/test_cosine_window.py b/tests/templates/test_state_preparations/test_cosine_window.py
@@ -20,6 +20,7 @@
 import pennylane as qml
 from pennylane.wires import WireError
 
+
 class TestDecomposition:
     """Tests that the template defines the correct decomposition."""