fix: add measure qubit targets in braket_program_context

amazon-braket · Apr 10, 2024 · 0bf0bb7 · 0bf0bb7
1 parent f556567
commit 0bf0bb7
Show file tree

Hide file tree

Showing 2 changed files with 7 additions and 7 deletions.
diff --git a/src/braket/circuits/braket_program_context.py b/src/braket/circuits/braket_program_context.py
@@ -170,3 +170,7 @@ def add_measure(self, target: tuple[int]) -> None:
         for index, qubit in enumerate(target):
             instruction = Instruction(Measure(index=index), qubit)
             self._circuit.add_instruction(instruction)
+            if self._circuit._measure_targets:
+                self._circuit._measure_targets.append(qubit)
+            else:
+                self._circuit._measure_targets = [qubit]
diff --git a/test/unit_tests/braket/circuits/test_circuit.py b/test/unit_tests/braket/circuits/test_circuit.py
@@ -839,18 +839,14 @@ def test_from_ir_round_trip_transformation():
                 "qubit[2] q;",
                 "h q[0];",
                 "cnot q[0], q[1];",
-                "b[0] = measure q[0];",
-                "b[1] = measure q[1];",
+                "b = measure q;",
             ]
         ),
         inputs={},
     )
-    new_ir = circuit.to_ir("OPENQASM")
-    new_circuit = Circuit.from_ir(new_ir)
 
-    assert new_ir == ir
-    assert Circuit.from_ir(source=ir.source, inputs=ir.inputs) == circuit
-    assert new_circuit == circuit
+    assert Circuit.from_ir(ir) == Circuit.from_ir(circuit.to_ir("OPENQASM"))
+    assert circuit.to_ir("OPENQASM") == Circuit.from_ir(ir).to_ir("OPENQASM")
 
 
 def test_add_with_instruction_with_default(cnot_instr):