allow operations after a measure if the target is not measured

src/braket/circuits/circuit.py

@@ -416,6 +416,44 @@ def _add_to_qubit_observable_set(self, result_type: ResultType) -> None:
         if isinstance(result_type, ObservableResultType) and result_type.target:
             self._qubit_observable_set.update(result_type.target)
 
+    def _check_if_qubit_measured(
+        self,
+        instruction: Instruction,
+        target: QubitSetInput | None = None,
+        target_mapping: dict[QubitInput, QubitInput] | None = None,
+    ) -> None:
+        """Checks if the target qubits are measured. If the qubit is already measured
+        the instruction will not be added to the Circuit.
+
+        Args:
+            instruction (Instruction): `Instruction` to add into `self`.
+            target (QubitSetInput | None): Target qubits for the
+                `instruction`. If a single qubit gate, an instruction is created for every index
+                in `target`.
+                Default = `None`.
+            target_mapping (dict[QubitInput, QubitInput] | None): A dictionary of
+                qubit mappings to apply to the `instruction.target`. Key is the qubit in
+                `instruction.target` and the value is what the key will be changed to.
+                Default = `None`.
+
+        Raises:
+            ValueError: If adding a gate or noise operation after a measure instruction.
+        """
+        print(self._measure_targets)
+        print("Instruction: ", instruction.target)
+        if (
+            target
+            and target in self._measure_targets
+            or (target_mapping and all(targ in self._measure_targets for targ in target_mapping))
+            or (
+                instruction.target
+                and all(targ in self._measure_targets for targ in instruction.target)
+            )
+        ):
+            raise ValueError(
+                "cannot add a gate or noise operation on a qubit after a measure instruction."
+            )
+
     def add_instruction(
         self,
         instruction: Instruction,
@@ -469,10 +507,8 @@ def add_instruction(
             raise TypeError("Only one of 'target_mapping' or 'target' can be supplied.")
 
         # Check if there is a measure instruction on the circuit
-        if not isinstance(instruction.operator, Measure) and any(
-            isinstance(instruction.operator, Measure) for instruction in self.instructions
-        ):
-            raise ValueError("cannot add a gate or noise after a measure instruction.")
+        if not isinstance(instruction.operator, Measure) and self._measure_targets:
+            self._check_if_qubit_measured(instruction, target, target_mapping)
 
         if not target_mapping and not target:
             # Nothing has been supplied, add instruction

test/unit_tests/braket/circuits/test_ascii_circuit_diagram.py

@@ -904,3 +904,34 @@ def test_measure_multiple_targets():
         "T  : |0|1|2|3|4|",
     )
     _assert_correct_diagram(circ, expected)
+
+
+def test_measure_multiple_instructions_after():
+    circ = (
+        Circuit()
+        .h(0)
+        .cnot(0, 1)
+        .cnot(1, 2)
+        .cnot(2, 3)
+        .measure(0)
+        .measure(1)
+        .h(3)
+        .cnot(3, 4)
+        .measure([2, 3])
+    )
+    expected = (
+        "T  : |0|1|2|3|4|5|6|",
+        "                    ",
+        "q0 : -H-C-----M-----",
+        "        |           ",
+        "q1 : ---X-C---M-----",
+        "          |         ",
+        "q2 : -----X-C-----M-",
+        "            |       ",
+        "q3 : -------X-H-C-M-",
+        "                |   ",
+        "q4 : -----------X---",
+        "",
+        "T  : |0|1|2|3|4|5|6|",
+    )
+    _assert_correct_diagram(circ, expected)

test/unit_tests/braket/circuits/test_circuit.py

@@ -725,9 +725,38 @@ def test_measure_empty_measure_after_measure_with_targets():
 
 
 def test_measure_gate_after():
-    message = "cannot add a gate or noise after a measure instruction."
+    message = "cannot add a gate or noise operation on a qubit after a measure instruction."
     with pytest.raises(ValueError, match=message):
-        Circuit().h(0).measure(0).h(1)
+        Circuit().h(0).measure(0).h(0)
+
+
+def test_measure_gate_after_with_target_mapping():
+    message = "cannot add a gate or noise operation on a qubit after a measure instruction."
+    instr = Instruction(Gate.CNot(), [0, 1])
+    with pytest.raises(ValueError, match=message):
+        Circuit().h(0).cnot(0, 1).cnot(1, 2).measure([0, 1]).add_instruction(
+            instr, target_mapping={0: 10, 1: 11}
+        )
+
+
+def test_measure_gate_after_with_target():
+    message = "cannot add a gate or noise operation on a qubit after a measure instruction."
+    instr = Instruction(Gate.CNot(), [0, 1])
+    with pytest.raises(ValueError, match=message):
+        Circuit().h(0).cnot(0, 1).cnot(1, 2).measure([0, 1]).add_instruction(instr, target=[10, 11])
+
+
+def test_measure_gate_after_measurement():
+    circ = Circuit().h(0).cnot(0, 1).cnot(1, 2).measure(0).h(2)
+    expected = (
+        Circuit()
+        .add_instruction(Instruction(Gate.H(), 0))
+        .add_instruction(Instruction(Gate.CNot(), [0, 1]))
+        .add_instruction(Instruction(Gate.CNot(), [1, 2]))
+        .add_instruction(Instruction(Measure(), 0))
+        .add_instruction(Instruction(Gate.H(), 2))
+    )
+    assert circ == expected
 
 
 def test_to_ir_with_measure():

test/unit_tests/braket/circuits/test_unicode_circuit_diagram.py

@@ -1062,3 +1062,38 @@ def test_measure_with_multiple_measures():
     )
     _assert_correct_diagram(circ, expected)
     _assert_correct_diagram(circ, expected)
+
+
+def test_measure_multiple_instructions_after():
+    circ = (
+        Circuit()
+        .h(0)
+        .cnot(0, 1)
+        .cnot(1, 2)
+        .cnot(2, 3)
+        .measure(0)
+        .measure(1)
+        .h(3)
+        .cnot(3, 4)
+        .measure([2, 3])
+    )
+    expected = (
+        "T  : │  0  │  1  │  2  │  3  │  4  │  5  │  6  │",
+        "      ┌───┐                   ┌───┐             ",
+        "q0 : ─┤ H ├───●───────────────┤ M ├─────────────",
+        "      └───┘   │               └───┘             ",
+        "            ┌─┴─┐             ┌───┐             ",
+        "q1 : ───────┤ X ├───●─────────┤ M ├─────────────",
+        "            └───┘   │         └───┘             ",
+        "                  ┌─┴─┐                   ┌───┐ ",
+        "q2 : ─────────────┤ X ├───●───────────────┤ M ├─",
+        "                  └───┘   │               └───┘ ",
+        "                        ┌─┴─┐ ┌───┐       ┌───┐ ",
+        "q3 : ───────────────────┤ X ├─┤ H ├───●───┤ M ├─",
+        "                        └───┘ └───┘   │   └───┘ ",
+        "                                    ┌─┴─┐       ",
+        "q4 : ───────────────────────────────┤ X ├───────",
+        "                                    └───┘       ",
+        "T  : │  0  │  1  │  2  │  3  │  4  │  5  │  6  │",
+    )
+    _assert_correct_diagram(circ, expected)