diff --git a/src/braket/circuits/braket_program_context.py b/src/braket/circuits/braket_program_context.py index 4371637d3..558692510 100644 --- a/src/braket/circuits/braket_program_context.py +++ b/src/braket/circuits/braket_program_context.py @@ -11,6 +11,7 @@ # ANY KIND, either express or implied. See the License for the specific # language governing permissions and limitations under the License. +from collections.abc import Iterable from typing import Optional, Union import numpy as np @@ -161,16 +162,17 @@ def handle_parameter_value( return FreeParameterExpression(evaluated_value) return value - def add_measure(self, target: tuple[int]) -> None: + def add_measure( + self, target: tuple[int], classical_targets: Optional[Iterable[int]] = None + ) -> None: """Add a measure instruction to the circuit Args: target (tuple[int]): the target qubits to be measured. + classical_targets (Optional[Iterable[int]]): the classical registers + to use in the qubit measurement. """ - for index, qubit in enumerate(target): + for iter, qubit in enumerate(target): + index = classical_targets[iter] if classical_targets else iter 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/src/braket/circuits/circuit.py b/src/braket/circuits/circuit.py index f15e03647..9ce7df7f8 100644 --- a/src/braket/circuits/circuit.py +++ b/src/braket/circuits/circuit.py @@ -504,6 +504,11 @@ def add_instruction( # Check if there is a measure instruction on the circuit self._check_if_qubit_measured(instruction, target, target_mapping) + # Update measure targets if instruction is a measurement + if isinstance(instruction.operator, Measure): + measure_target = target or instruction.target[0] + self._measure_targets = (self._measure_targets or []) + [measure_target] + if not target_mapping and not target: # Nothing has been supplied, add instruction instructions_to_add = [instruction] @@ -710,10 +715,6 @@ def _add_measure(self, target_qubits: QubitSetInput) -> None: target=target, ) ) - if self._measure_targets: - self._measure_targets.append(target) - else: - self._measure_targets = [target] def measure(self, target_qubits: QubitSetInput) -> Circuit: """