Use qubit types from braket.registers module (#771)

src/braket/experimental/autoqasm/instructions/qubits.py

@@ -22,8 +22,11 @@
 from openpulse.printer import dumps
 
 from braket.experimental.autoqasm import constants, errors, program
+from braket.registers.qubit import Qubit
 
-QubitIdentifierType = Union[int, oqpy._ClassicalVar, oqpy.base.OQPyExpression, str, oqpy.Qubit]
+QubitIdentifierType = Union[
+    int, str, Qubit, oqpy._ClassicalVar, oqpy.base.OQPyExpression, oqpy.Qubit
+]
 
 
 def is_qubit_identifier_type(qubit: Any) -> bool:
@@ -58,11 +61,8 @@ def _get_physical_qubit_indices(qids: list[str]) -> list[int]:
     return braket_qubits
 
 
-def _global_qubit_register(qubit_idx_expr: Union[int, str]) -> str:
-    # TODO: We should index into a oqpy.Qubit register rather
-    # than manually generating the string to index into
-    # a hard-coded global qubit register.
-    return f"{constants.QUBIT_REGISTER}[{qubit_idx_expr}]"
+def _global_qubit_register(qubit_idx_expr: Union[int, str]) -> oqpy.Qubit:
+    return oqpy.Qubit(f"{constants.QUBIT_REGISTER}[{qubit_idx_expr}]", needs_declaration=False)
 
 
 @singledispatch
@@ -92,15 +92,15 @@ def _(qid: float) -> oqpy.Qubit:
 def _(qid: int) -> oqpy.Qubit:
     # Integer virtual qubit, like `h(0)`
     program.get_program_conversion_context().register_qubit(qid)
-    return oqpy.Qubit(_global_qubit_register(qid), needs_declaration=False)
+    return _global_qubit_register(qid)
 
 
 @_qubit.register
 def _(qid: oqpy._ClassicalVar) -> oqpy.Qubit:
     # Indexed by variable, such as i in range(n); h(i)
     if program.get_program_conversion_context().get_declared_qubits() is None:
         raise errors.UnknownQubitCountError()
-    return oqpy.Qubit(_global_qubit_register(qid.name), needs_declaration=False)
+    return _global_qubit_register(qid.name)
 
 
 @_qubit.register
@@ -111,7 +111,7 @@ def _(qid: oqpy.base.OQPyExpression) -> oqpy.Qubit:
 
     oqpy_program = program.get_program_conversion_context().get_oqpy_program()
     qubit_idx_expr = dumps(qid.to_ast(oqpy_program))
-    return oqpy.Qubit(_global_qubit_register(qubit_idx_expr), needs_declaration=False)
+    return _global_qubit_register(qubit_idx_expr)
 
 
 @_qubit.register

src/braket/experimental/autoqasm/pulse/pulse.py

@@ -19,7 +19,6 @@
 
 import oqpy
 
-from braket.circuits.qubit_set import QubitSet
 from braket.experimental.autoqasm import program as aq_program
 from braket.experimental.autoqasm.instructions.qubits import (
     QubitIdentifierType,
@@ -32,6 +31,7 @@
 from braket.pulse.frame import Frame
 from braket.pulse.pulse_sequence import _validate_uniqueness
 from braket.pulse.waveforms import Waveform
+from braket.registers.qubit_set import QubitSet
 
 
 def _pulse_instruction(name: str, frame: Frame, *args) -> None: