diff --git a/src/braket/circuits/ascii_circuit_diagram.py b/src/braket/circuits/ascii_circuit_diagram.py index 8e7e97421..81a8d1f53 100644 --- a/src/braket/circuits/ascii_circuit_diagram.py +++ b/src/braket/circuits/ascii_circuit_diagram.py @@ -140,8 +140,11 @@ def _ascii_group_items( ): continue - # TODO: We should use isinstance(item.operator, GPhase) - if isinstance(item, Instruction) and item.operator.__class__.__name__ == "GPhase": + if ( + isinstance(item, Instruction) + and isinstance(item.operator, Gate) + and item.operator.name == "GPhase" + ): qubit_range = QubitSet() elif (isinstance(item, ResultType) and not item.target) or ( isinstance(item, Instruction) and isinstance(item.operator, CompilerDirective) @@ -162,7 +165,9 @@ def _ascii_group_items( instr_group = group[1] # Take into account overlapping multi-qubit gates if not qubits_added.intersection(set(qubit_range)) or ( - isinstance(item, Instruction) and item.operator.__class__.__name__ == "GPhase" + isinstance(item, Instruction) + and isinstance(item.operator, Gate) + and item.operator.name == "GPhase" ): instr_group.append(item) qubits_added.update(qubit_range) @@ -294,7 +299,11 @@ def _ascii_diagram_column( num_after = len(circuit_qubits) - 1 after = ["|"] * (num_after - 1) + ([marker] if num_after else []) ascii_symbols = [ascii_symbol] + after - elif isinstance(item, Instruction) and item.operator.__class__.__name__ == "GPhase": + elif ( + isinstance(item, Instruction) + and isinstance(item.operator, Gate) + and item.operator.name == "GPhase" + ): target_qubits = circuit_qubits control_qubits = QubitSet() target_and_control = QubitSet() diff --git a/src/braket/circuits/moments.py b/src/braket/circuits/moments.py index ddf608162..6e87db78d 100644 --- a/src/braket/circuits/moments.py +++ b/src/braket/circuits/moments.py @@ -19,6 +19,7 @@ from typing import Any, NamedTuple, Union from braket.circuits.compiler_directive import CompilerDirective +from braket.circuits.gate import Gate from braket.circuits.instruction import Instruction from braket.circuits.noise import Noise from braket.registers.qubit import Qubit @@ -184,7 +185,7 @@ def _add(self, instruction: Instruction, noise_index: int = 0) -> None: self._time_all_qubits = time elif isinstance(operator, Noise): self.add_noise(instruction) - elif operator.__class__.__name__ == "GPhase": + elif isinstance(operator, Gate) and operator.name == "GPhase": time = self._get_qubit_times(self._max_times.keys()) + 1 self._number_gphase_in_current_moment += 1 key = MomentsKey(