fix: Refactor Qubit and QubitSet to a separate module

src/braket/circuits/ascii_circuit_diagram.py

@@ -22,8 +22,8 @@
 from braket.circuits.gate import Gate
 from braket.circuits.instruction import Instruction
 from braket.circuits.noise import Noise
-from braket.circuits.qubit_set import QubitSet
 from braket.circuits.result_type import ResultType
+from braket.registers.qubit_set import QubitSet
 
 
 class AsciiCircuitDiagram(CircuitDiagram):

src/braket/circuits/circuit.py

@@ -39,8 +39,6 @@
 from braket.circuits.observable import Observable
 from braket.circuits.observables import TensorProduct
 from braket.circuits.parameterizable import Parameterizable
-from braket.circuits.qubit import QubitInput
-from braket.circuits.qubit_set import QubitSet, QubitSetInput
 from braket.circuits.result_type import (
     ObservableParameterResultType,
     ObservableResultType,
@@ -60,6 +58,8 @@
 from braket.pulse import ArbitraryWaveform, Frame
 from braket.pulse.ast.qasm_parser import ast_to_qasm
 from braket.pulse.pulse_sequence import PulseSequence, _validate_uniqueness
+from braket.registers.qubit import QubitInput
+from braket.registers.qubit_set import QubitSet, QubitSetInput
 
 SubroutineReturn = TypeVar(
     "SubroutineReturn", Iterable[Instruction], Instruction, ResultType, Iterable[ResultType]

src/braket/circuits/compiler_directive.py

@@ -16,8 +16,8 @@
 from typing import Any, Sequence, Tuple
 
 from braket.circuits.operator import Operator
-from braket.circuits.qubit_set import QubitSet
 from braket.circuits.serialization import IRType, SerializationProperties
+from braket.registers.qubit_set import QubitSet
 
 
 class CompilerDirective(Operator):

src/braket/circuits/gate.py

@@ -18,12 +18,12 @@
 
 from braket.circuits.basis_state import BasisState, BasisStateInput
 from braket.circuits.quantum_operator import QuantumOperator
-from braket.circuits.qubit_set import QubitSet
 from braket.circuits.serialization import (
     IRType,
     OpenQASMSerializationProperties,
     SerializationProperties,
 )
+from braket.registers.qubit_set import QubitSet
 
 
 class Gate(QuantumOperator):

src/braket/circuits/gate_calibrations.py

@@ -17,13 +17,13 @@
 from typing import Any, Dict, List, Optional, Tuple
 
 from braket.circuits.gate import Gate
-from braket.circuits.qubit_set import QubitSet
 from braket.circuits.serialization import (
     IRType,
     OpenQASMSerializationProperties,
     QubitReferenceType,
 )
 from braket.pulse.pulse_sequence import PulseSequence
+from braket.registers.qubit_set import QubitSet
 
 
 class GateCalibrations:

src/braket/circuits/gates.py

@@ -39,11 +39,11 @@
     is_unitary,
     verify_quantum_operator_matrix_dimensions,
 )
-from braket.circuits.qubit import QubitInput
-from braket.circuits.qubit_set import QubitSet, QubitSetInput
 from braket.circuits.serialization import OpenQASMSerializationProperties
 from braket.pulse.ast.qasm_parser import ast_to_qasm
 from braket.pulse.pulse_sequence import PulseSequence
+from braket.registers.qubit import QubitInput
+from braket.registers.qubit_set import QubitSet, QubitSetInput
 
 """
 To add a new gate:

src/braket/circuits/instruction.py

@@ -20,9 +20,9 @@
 from braket.circuits.gate import Gate
 from braket.circuits.operator import Operator
 from braket.circuits.quantum_operator import QuantumOperator
-from braket.circuits.qubit import QubitInput
-from braket.circuits.qubit_set import QubitSet, QubitSetInput
 from braket.circuits.serialization import IRType, SerializationProperties
+from braket.registers.qubit import QubitInput
+from braket.registers.qubit_set import QubitSet, QubitSetInput
 
 # InstructionOperator is a type alias, and it can be expanded to include other operators
 InstructionOperator = Operator

src/braket/circuits/moments.py

@@ -31,8 +31,8 @@
 from braket.circuits.compiler_directive import CompilerDirective
 from braket.circuits.instruction import Instruction
 from braket.circuits.noise import Noise
-from braket.circuits.qubit import Qubit
-from braket.circuits.qubit_set import QubitSet
+from braket.registers.qubit import Qubit
+from braket.registers.qubit_set import QubitSet
 
 
 class MomentType(str, Enum):

src/braket/circuits/noise.py

@@ -21,12 +21,12 @@
 from braket.circuits.free_parameter_expression import FreeParameterExpression
 from braket.circuits.parameterizable import Parameterizable
 from braket.circuits.quantum_operator import QuantumOperator
-from braket.circuits.qubit_set import QubitSet
 from braket.circuits.serialization import (
     IRType,
     OpenQASMSerializationProperties,
     SerializationProperties,
 )
+from braket.registers.qubit_set import QubitSet
 
 
 class Noise(QuantumOperator):

src/braket/circuits/noise_helpers.py

@@ -23,7 +23,7 @@
 from braket.circuits.moments import Moments
 from braket.circuits.noise import Noise
 from braket.circuits.quantum_operator_helpers import is_unitary
-from braket.circuits.qubit_set import QubitSet, QubitSetInput
+from braket.registers.qubit_set import QubitSet, QubitSetInput
 
 if TYPE_CHECKING:  # pragma: no cover
     from braket.circuits.circuit import Circuit

src/braket/circuits/noise_model/circuit_instruction_criteria.py

@@ -16,7 +16,7 @@
 
 from braket.circuits.instruction import Instruction
 from braket.circuits.noise_model.criteria import Criteria
-from braket.circuits.qubit_set import QubitSetInput
+from braket.registers.qubit_set import QubitSetInput
 
 
 class CircuitInstructionCriteria(Criteria):

src/braket/circuits/noise_model/criteria_input_parsing.py

@@ -14,7 +14,7 @@
 from typing import Iterable, Optional, Set, Tuple, Union
 
 from braket.circuits.quantum_operator import QuantumOperator
-from braket.circuits.qubit_set import QubitSetInput
+from braket.registers.qubit_set import QubitSetInput
 
 
 def parse_operator_input(

src/braket/circuits/noise_model/gate_criteria.py

@@ -21,7 +21,7 @@
     parse_operator_input,
     parse_qubit_input,
 )
-from braket.circuits.qubit_set import QubitSetInput
+from braket.registers.qubit_set import QubitSetInput
 
 
 class GateCriteria(CircuitInstructionCriteria):

src/braket/circuits/noise_model/initialization_criteria.py

@@ -14,7 +14,7 @@
 from abc import abstractmethod
 
 from braket.circuits.noise_model.criteria import Criteria
-from braket.circuits.qubit_set import QubitSetInput
+from braket.registers.qubit_set import QubitSetInput
 
 
 class InitializationCriteria(Criteria):

src/braket/circuits/noise_model/noise_model.py

@@ -25,8 +25,8 @@
 from braket.circuits.noise_model.criteria import Criteria, CriteriaKey, CriteriaKeyResult
 from braket.circuits.noise_model.initialization_criteria import InitializationCriteria
 from braket.circuits.noise_model.result_type_criteria import ResultTypeCriteria
-from braket.circuits.qubit_set import QubitSetInput
 from braket.circuits.result_types import ObservableResultType
+from braket.registers.qubit_set import QubitSetInput
 
 
 @dataclass

src/braket/circuits/noise_model/observable_criteria.py

@@ -20,8 +20,8 @@
 )
 from braket.circuits.noise_model.result_type_criteria import ResultTypeCriteria
 from braket.circuits.observable import Observable
-from braket.circuits.qubit_set import QubitSetInput
 from braket.circuits.result_type import ObservableResultType, ResultType
+from braket.registers.qubit_set import QubitSetInput
 
 
 class ObservableCriteria(ResultTypeCriteria):

src/braket/circuits/noise_model/qubit_initialization_criteria.py

@@ -16,7 +16,7 @@
 from braket.circuits.noise_model.criteria import Criteria, CriteriaKey, CriteriaKeyResult
 from braket.circuits.noise_model.criteria_input_parsing import parse_qubit_input
 from braket.circuits.noise_model.initialization_criteria import InitializationCriteria
-from braket.circuits.qubit_set import QubitSet, QubitSetInput
+from braket.registers.qubit_set import QubitSet, QubitSetInput
 
 
 class QubitInitializationCriteria(InitializationCriteria):

src/braket/circuits/noise_model/unitary_gate_criteria.py

@@ -18,7 +18,7 @@
 from braket.circuits.noise_model.circuit_instruction_criteria import CircuitInstructionCriteria
 from braket.circuits.noise_model.criteria import Criteria, CriteriaKey, CriteriaKeyResult
 from braket.circuits.noise_model.criteria_input_parsing import parse_qubit_input
-from braket.circuits.qubit_set import QubitSetInput
+from braket.registers.qubit_set import QubitSetInput
 
 
 class UnitaryGateCriteria(CircuitInstructionCriteria):

src/braket/circuits/noises.py

@@ -36,9 +36,9 @@
     is_cptp,
     verify_quantum_operator_matrix_dimensions,
 )
-from braket.circuits.qubit import QubitInput
-from braket.circuits.qubit_set import QubitSet, QubitSetInput
 from braket.circuits.serialization import OpenQASMSerializationProperties
+from braket.registers.qubit import QubitInput
+from braket.registers.qubit_set import QubitSet, QubitSetInput
 
 """
 To add a new Noise implementation:

src/braket/circuits/observable.py

@@ -21,12 +21,12 @@
 
 from braket.circuits.gate import Gate
 from braket.circuits.quantum_operator import QuantumOperator
-from braket.circuits.qubit_set import QubitSet
 from braket.circuits.serialization import (
     IRType,
     OpenQASMSerializationProperties,
     SerializationProperties,
 )
+from braket.registers.qubit_set import QubitSet
 
 
 class Observable(QuantumOperator):

src/braket/circuits/observables.py

@@ -29,8 +29,8 @@
     is_hermitian,
     verify_quantum_operator_matrix_dimensions,
 )
-from braket.circuits.qubit_set import QubitSet
 from braket.circuits.serialization import IRType, OpenQASMSerializationProperties
+from braket.registers.qubit_set import QubitSet
 
 
 class H(StandardObservable):

src/braket/circuits/qubit.py

@@ -11,58 +11,4 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 
-from __future__ import annotations
-
-import numbers
-from typing import Union
-
-QubitInput = Union["Qubit", int]
-
-
-class Qubit(int):
-    """
-    A quantum bit index. The index of this qubit is locally scoped towards the contained
-    circuit. This may not be the exact qubit index on the quantum device.
-    """
-
-    def __new__(cls, index: int):
-        """
-        Args:
-            index (int): Index of the qubit.
-
-        Raises:
-            ValueError: If `index` is less than zero.
-
-        Examples:
-            >>> Qubit(0)
-            >>> Qubit(1)
-        """
-        if not isinstance(index, numbers.Integral):
-            raise TypeError(f"Supplied qubit index, {index}, must be an integer.")
-        if index < 0:
-            raise ValueError(f"Supplied qubit index, {index}, cannot be less than zero.")
-        return super().__new__(cls, index)
-
-    def __repr__(self):
-        return f"Qubit({super().__repr__()})"
-
-    def __str__(self):
-        return self.__repr__()
-
-    @staticmethod
-    def new(qubit: QubitInput) -> Qubit:
-        """
-        Helper constructor - if input is a `Qubit` it returns the same value,
-        else a new `Qubit` is constructed.
-
-        Args:
-            qubit (QubitInput): `Qubit` index. If `type == Qubit` then the `qubit` is returned.
-
-        Returns:
-            Qubit: The qubit.
-        """
-
-        if isinstance(qubit, Qubit):
-            return qubit
-        else:
-            return Qubit(qubit)
+from braket.registers import Qubit, QubitInput  # noqa: F401

src/braket/circuits/qubit_set.py

@@ -11,83 +11,4 @@
 # ANY KIND, either express or implied. See the License for the specific
 # language governing permissions and limitations under the License.
 
-from __future__ import annotations
-
-from typing import Any, Dict, Iterable, Union
-
-from boltons.setutils import IndexedSet
-
-from braket.circuits.qubit import Qubit, QubitInput
-
-QubitSetInput = Union[QubitInput, Iterable[QubitInput]]
-
-
-def _flatten(other: Any) -> Any:
-    if isinstance(other, Iterable) and not isinstance(other, str):
-        for item in other:
-            yield from _flatten(item)
-    else:
-        yield other
-
-
-class QubitSet(IndexedSet):
-    """
-    An ordered, unique set of quantum bits.
-
-    Note:
-        QubitSet implements `__hash__()` but is a mutable object, therefore be careful when
-        mutating this object.
-    """
-
-    def __init__(self, qubits: QubitSetInput = None):
-        """
-        Args:
-            qubits (QubitSetInput): Qubits to be included in the `QubitSet`. Default is `None`.
-
-        Examples:
-            >>> qubits = QubitSet([0, 1])
-            >>> for qubit in qubits:
-            ...     print(qubit)
-            ...
-            Qubit(0)
-            Qubit(1)
-
-            >>> qubits = QubitSet([0, 1, [2, 3]])
-            >>> for qubit in qubits:
-            ...     print(qubit)
-            ...
-            Qubit(0)
-            Qubit(1)
-            Qubit(2)
-            Qubit(3)
-        """
-
-        _qubits = [Qubit.new(qubit) for qubit in _flatten(qubits)] if qubits is not None else None
-        super().__init__(_qubits)
-
-    def map(self, mapping: Dict[QubitInput, QubitInput]) -> QubitSet:
-        """
-        Creates a new `QubitSet` where this instance's qubits are mapped to the values in `mapping`.
-        If this instance contains a qubit that is not in the `mapping` that qubit is not modified.
-
-        Args:
-            mapping (Dict[QubitInput, QubitInput]): A dictionary of qubit mappings to
-                apply. Key is the qubit in this instance to target, and the value is what
-                the key will be changed to.
-
-        Returns:
-            QubitSet: A new QubitSet with the `mapping` applied.
-
-        Examples:
-            >>> qubits = QubitSet([0, 1])
-            >>> mapping = {0: 10, Qubit(1): Qubit(11)}
-            >>> qubits.map(mapping)
-            QubitSet([Qubit(10), Qubit(11)])
-        """
-
-        new_qubits = [mapping.get(qubit, qubit) for qubit in self]
-
-        return QubitSet(new_qubits)
-
-    def __hash__(self):
-        return hash(tuple(self))
+from braket.registers.qubit_set import QubitSet, QubitSetInput  # noqa: F401

src/braket/circuits/result_type.py

@@ -18,13 +18,13 @@
 from braket.circuits.free_parameter import FreeParameter
 from braket.circuits.observable import Observable
 from braket.circuits.observables import Sum
-from braket.circuits.qubit import QubitInput
-from braket.circuits.qubit_set import QubitSet, QubitSetInput
 from braket.circuits.serialization import (
     IRType,
     OpenQASMSerializationProperties,
     SerializationProperties,
 )
+from braket.registers.qubit import QubitInput
+from braket.registers.qubit_set import QubitSet, QubitSetInput
 
 
 class ResultType: