Repository: PennyLaneAI/pennylane
Pull request #7404 (2025-05-09 18:01:57)

Added Qubit Manager class
**Context:**
Added QubitManager class for tracking auxiliary qubits.

**Description of the Change:**

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**


[sc-90641]


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 4ad7ed5f9..45b199141 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -406,6 +406,8 @@
   [(#7229)](https://github.com/PennyLaneAI/pennylane/pull/7229)
   [(#7333)](https://github.com/PennyLaneAI/pennylane/pull/7333)
 
+* `pennylane.labs.QubitManager`, `pennylane.labs.AllocWires`, and `pennylane.labs.FreeWires` classes have been added to track and manage auxilliary qubits.
+  [(#7404)](https://github.com/PennyLaneAI/pennylane/pull/7404)
 
 <h3>Breaking changes 💔</h3>
 
@@ -627,6 +629,7 @@ This release contains contributions from (in alphabetical order):
 Guillermo Alonso-Linaje,
 Astral Cai,
 Yushao Chen,
+Diksha Dhawan,
 Marcus Edwards,
 Lillian Frederiksen,
 Pietropaolo Frisoni,
diff --git a/pennylane/labs/resource_estimation/__init__.py b/pennylane/labs/resource_estimation/__init__.py
index fc4f3a7f1..f29b824d5 100644
--- a/pennylane/labs/resource_estimation/__init__.py
+++ b/pennylane/labs/resource_estimation/__init__.py
@@ -22,4 +22,16 @@ resource estimation.
 
 .. currentmodule:: pennylane.labs.resource_estimation
 
+Qubit Management Classes:
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autosummary::
+    :toctree: api
+
+    ~QubitManager
+    ~AllocWires
+    ~FreeWires
+
 """
+
+from .qubit_manager import AllocWires, FreeWires, QubitManager
diff --git a/pennylane/labs/resource_estimation/qubit_manager.py b/pennylane/labs/resource_estimation/qubit_manager.py
new file mode 100644
index 000000000..eaa422689
--- /dev/null
+++ b/pennylane/labs/resource_estimation/qubit_manager.py
@@ -0,0 +1,196 @@
+# Copyright 2025 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+r"""This module contains the base class for qubit management"""
+
+from typing import Union
+
+import pennylane as qml
+
+# pylint: disable= too-few-public-methods
+
+
+class QubitManager:
+    r"""Manages and tracks the auxiliary and algorithmic qubits used in a quantum circuit.
+
+    Args:
+        work_wires (int or Dict[str, int]): Number of work wires or a dictionary containing
+            number of clean and dirty work wires. All ``work_wires`` are assumed to be clean when
+            `int` is provided.
+        algo_wires (int): Number of algorithmic wires, default value is ``0``.
+        tight_budget (bool): Determines whether extra clean qubits can be allocated when they
+            exceed the available amount. The default is ``False``.
+
+    **Example**
+
+    >>> q = QubitManager(
+    ...             work_wires={"clean": 2, "dirty": 2},
+    ...             tight_budget=False,
+    ...     )
+    >>> print(q)
+    QubitManager(clean=2, dirty=2, logic=0, tight_budget=False)
+
+    """
+
+    def __init__(self, work_wires: Union[int, dict], algo_wires=0, tight_budget=False) -> None:
+
+        if isinstance(work_wires, dict):
+            clean_wires = work_wires["clean"]
+            dirty_wires = work_wires["dirty"]
+        else:
+            clean_wires = work_wires
+            dirty_wires = 0
+
+        self.tight_budget = tight_budget
+        self._logic_qubit_counts = algo_wires
+        self._clean_qubit_counts = clean_wires
+        self._dirty_qubit_counts = dirty_wires
+
+    def __str__(self):
+        return (
+            f"QubitManager(clean qubits={self._clean_qubit_counts}, dirty qubits={self._dirty_qubit_counts}, "
+            f"algorithmic qubits={self._logic_qubit_counts}, tight budget={self.tight_budget})"
+        )
+
+    def __repr__(self) -> str:
+        work_wires_str = repr(
+            {"clean": self._clean_qubit_counts, "dirty": self._dirty_qubit_counts}
+        )
+        return (
+            f"QubitManager(work_wires={work_wires_str}, algo_wires={self._logic_qubit_counts}, "
+            f"tight_budget={self.tight_budget})"
+        )
+
+    def __eq__(self, other: object) -> bool:
+        return (
+            isinstance(other, self.__class__)
+            and (self.clean_qubits == other.clean_qubits)
+            and (self.dirty_qubits == other.dirty_qubits)
+            and (self.algo_qubits == other.algo_qubits)
+            and (self.tight_budget == other.tight_budget)
+        )
+
+    @property
+    def clean_qubits(self):
+        r"""Returns the number of clean qubits."""
+        return self._clean_qubit_counts
+
+    @property
+    def dirty_qubits(self):
+        r"""Returns the number of dirty qubits."""
+        return self._dirty_qubit_counts
+
+    @property
+    def algo_qubits(self):
+        r"""Returns the number of algorithmic qubits."""
+        return self._logic_qubit_counts
+
+    @property
+    def total_qubits(self):
+        r"""Returns the number of total qubits."""
+        return self._clean_qubit_counts + self._dirty_qubit_counts + self.algo_qubits
+
+    @algo_qubits.setter
+    def algo_qubits(self, count: int):  # these get set manually, the rest are dynamically updated
+        r"""Setter for algorithmic qubits."""
+        self._logic_qubit_counts = count
+
+    def allocate_qubits(self, num_qubits: int):
+        r"""Allocates extra clean qubits.
+
+        Args:
+            num_qubits(int): number of qubits to be allocated
+
+        """
+        self._clean_qubit_counts += num_qubits
+
+    def grab_clean_qubits(self, num_qubits: int):
+        r"""Grabs clean qubits.
+
+        Args:
+            num_qubits(int) : number of clean qubits to be grabbed
+
+        Raises:
+            ValueError: If tight_budget is `True` number of qubits to be grabbed is greater than
+            available clean qubits.
+
+        """
+        available_clean = self.clean_qubits
+
+        if num_qubits > available_clean:
+            if self.tight_budget:
+                raise ValueError(
+                    f"Grabbing more qubits than available clean qubits."
+                    f"Number of clean qubits available is {available_clean}, while {num_qubits} are being grabbed."
+                )
+            self._clean_qubit_counts = 0
+        else:
+            self._clean_qubit_counts -= num_qubits
+        self._dirty_qubit_counts += num_qubits
+
+    def free_qubits(self, num_qubits: int):
+        r"""Frees dirty qubits and converts them to clean qubits.
+
+        Args:
+            num_qubits(int) : number of qubits to be freed
+
+        Raises:
+            ValueError: If number of qubits to be freed is greater than available dirty qubits.
+        """
+
+        if num_qubits > self.dirty_qubits:
+            raise ValueError(
+                f"Freeing more qubits than available dirty qubits."
+                f"Number of dirty qubits available is {self.dirty_qubits}, while {num_qubits} qubits are being released."
+            )
+
+        self._dirty_qubit_counts -= num_qubits
+        self._clean_qubit_counts += num_qubits
+
+
+class _WireAction:
+    """Base class for operations that manage qubit resources."""
+
+    _queue_category = "_resource_qubit_action"
+
+    def __init__(self, num_wires):
+        self.num_wires = num_wires
+        if qml.QueuingManager.recording():
+            self.queue()
+
+    def queue(self, context=qml.QueuingManager):
+        r"""Adds the wire action object to a queue."""
+        context.append(self)
+        return self
+
+
+class AllocWires(_WireAction):
+    r"""Allows users to allocate clean work wires.
+
+    Args:
+        num_wires (int): number of work wires to be allocated.
+    """
+
+    def __repr__(self) -> str:
+        return f"AllocWires({self.num_wires})"
+
+
+class FreeWires(_WireAction):
+    r"""Allows users to free dirty work wires.
+
+    Args:
+        num_wires (int): number of dirty work wires to be freed.
+    """
+
+    def __repr__(self) -> str:
+        return f"FreeWires({self.num_wires})"


Test patch:
diff --git a/pennylane/labs/tests/resource_estimation/test_qubit_manager.py b/pennylane/labs/tests/resource_estimation/test_qubit_manager.py
new file mode 100644
index 000000000..8387b4fa1
--- /dev/null
+++ b/pennylane/labs/tests/resource_estimation/test_qubit_manager.py
@@ -0,0 +1,217 @@
+# Copyright 2018-2025 Xanadu Quantum Technologies Inc.
+
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+
+#     http://www.apache.org/licenses/LICENSE-2.0
+
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+This module contains tests for classes needed to track auxilliary qubits.
+"""
+import copy
+
+import pytest
+
+import pennylane as qml
+from pennylane.labs.resource_estimation import AllocWires, FreeWires, QubitManager
+
+
+# pylint: disable= no-self-use
+class TestQubitManager:
+    """Test the methods and attributes of the QubitManager class"""
+
+    qm_quantities = (
+        QubitManager(work_wires=2),
+        QubitManager(work_wires={"clean": 4, "dirty": 2}, algo_wires=20),
+        QubitManager({"clean": 2, "dirty": 2}, algo_wires=10, tight_budget=True),
+    )
+
+    qm_parameters = (
+        (2, 0, 0, False),
+        (4, 2, 20, False),
+        (2, 2, 10, True),
+    )
+
+    @pytest.mark.parametrize("qm, attribute_tup", zip(qm_quantities, qm_parameters))
+    def test_init(self, qm, attribute_tup):
+        """Test that the QubitManager class is instantiated as expected."""
+        clean_qubits, dirty_qubits, logic_qubits, tight_budget = attribute_tup
+
+        assert qm.clean_qubits == clean_qubits
+        assert qm.dirty_qubits == dirty_qubits
+        assert qm.algo_qubits == logic_qubits
+        assert qm.tight_budget == tight_budget
+
+    @pytest.mark.parametrize("qm, attribute_tup", zip(qm_quantities, qm_parameters))
+    def test_equality(self, qm, attribute_tup):
+        """Test that the equality methods behaves as expected"""
+
+        clean_qubits, dirty_qubits, algo_qubits, tight_budget = attribute_tup
+
+        qm2 = QubitManager(
+            work_wires={"clean": clean_qubits, "dirty": dirty_qubits},
+            algo_wires=algo_qubits,
+            tight_budget=tight_budget,
+        )
+        assert qm == qm2
+
+    extra_qubits = (0, 2, 4)
+
+    @pytest.mark.parametrize(
+        "qm, attribute_tup, alloc_q", zip(copy.deepcopy(qm_quantities), qm_parameters, extra_qubits)
+    )
+    def test_allocate_qubits(self, qm, attribute_tup, alloc_q):
+        """Test that the extra qubits are allocated correctly."""
+
+        clean_qubits = attribute_tup[0]
+
+        qm.allocate_qubits(alloc_q)
+        assert qm.clean_qubits == clean_qubits + alloc_q
+
+    qm_parameters_algo = (
+        (2, 0, 0, False),
+        (4, 2, 2, False),
+        (2, 2, 4, True),
+    )
+
+    @pytest.mark.parametrize(
+        "qm, attribute_tup, algo_q",
+        zip(copy.deepcopy(qm_quantities), qm_parameters_algo, extra_qubits),
+    )
+    def test_repr(self, qm, attribute_tup, algo_q):
+        """Test that the QubitManager representation is correct."""
+
+        clean_qubits, dirty_qubits, logic_qubits, tight_budget = attribute_tup
+
+        work_wires_str = repr({"clean": clean_qubits, "dirty": dirty_qubits})
+        expected_string = (
+            f"QubitManager(work_wires={work_wires_str}, algo_wires={logic_qubits}, "
+            f"tight_budget={tight_budget})"
+        )
+
+        qm.algo_qubits = algo_q
+        assert repr(qm) == expected_string
+
+    @pytest.mark.parametrize(
+        "qm, attribute_tup, algo_q",
+        zip(copy.deepcopy(qm_quantities), qm_parameters_algo, extra_qubits),
+    )
+    def test_str(self, qm, attribute_tup, algo_q):
+        """Test that the QubitManager string is correct."""
+
+        clean_qubits, dirty_qubits, logic_qubits, tight_budget = attribute_tup
+
+        expected_string = (
+            f"QubitManager(clean qubits={clean_qubits}, dirty qubits={dirty_qubits}, "
+            f"algorithmic qubits={logic_qubits}, tight budget={tight_budget})"
+        )
+        qm.algo_qubits = algo_q
+        assert str(qm) == expected_string
+
+    @pytest.mark.parametrize(
+        "qm, attribute_tup, algo_q",
+        zip(copy.deepcopy(qm_quantities), qm_parameters_algo, extra_qubits),
+    )
+    def test_algo_qubits(self, qm, attribute_tup, algo_q):
+        """Test that the logic qubits are set correctly."""
+
+        logic_qubits = attribute_tup[2]
+
+        qm.algo_qubits = algo_q
+        assert qm.algo_qubits == logic_qubits
+
+    @pytest.mark.parametrize(
+        "qm, attribute_tup, algo_q",
+        zip(copy.deepcopy(qm_quantities), qm_parameters_algo, extra_qubits),
+    )
+    def test_total_qubits(self, qm, attribute_tup, algo_q):
+        """Test that the total qubits returned are correct."""
+
+        qm.algo_qubits = algo_q
+        total_qubits = attribute_tup[0] + attribute_tup[1] + attribute_tup[2]
+        assert qm.total_qubits == total_qubits
+
+    def test_grab_clean_qubits(self):
+        """Test that the clean qubits are grabbed properly."""
+
+        qm = QubitManager(work_wires={"clean": 4, "dirty": 2}, tight_budget=False)
+        qm.grab_clean_qubits(6)
+        assert qm.clean_qubits == 0
+        assert qm.dirty_qubits == 8
+
+    def test_error_grab_clean_qubits(self):
+        """Test that an error is raised when the number of clean qubits required is greater
+        than the available qubits."""
+
+        qm = QubitManager(work_wires={"clean": 4, "dirty": 2}, tight_budget=True)
+        with pytest.raises(ValueError, match="Grabbing more qubits than available clean qubits."):
+            qm.grab_clean_qubits(6)
+
+    def test_free_qubits(self):
+        """Test that the dirty qubits are freed properly."""
+
+        qm = QubitManager(work_wires={"clean": 4, "dirty": 2})
+        qm.free_qubits(2)
+        assert qm.clean_qubits == 6
+        assert qm.dirty_qubits == 0
+
+    def test_error_free_qubits(self):
+        """Test that an error is raised when the number of qubits being freed is greater
+        than the available dirty qubits."""
+
+        qm = QubitManager(work_wires={"clean": 4, "dirty": 2}, tight_budget=True)
+        with pytest.raises(ValueError, match="Freeing more qubits than available dirty qubits."):
+            qm.free_qubits(6)
+
+
+class TestAllocWires:
+    """Test the methods and attributes of the AllocWires class"""
+
+    def test_init(self):
+        """Test that the AllocWires class is instantiated as expected when there is no active recording."""
+
+        for i in range(3):
+            assert AllocWires(i).num_wires == i
+
+    def test_repr(self):
+        """Test that correct representation is returned for AllocWires class"""
+
+        wires = 3
+        exp_string = f"AllocWires({wires})"
+        assert repr(AllocWires(wires)) == exp_string
+
+    def test_init_recording(self):
+        """Test that the AllocWires class is instantiated as expected when there is active recording."""
+        with qml.queuing.AnnotatedQueue() as q:
+            ops = [AllocWires(2), AllocWires(4)]
+        assert q.queue == ops
+
+
+class TestFreeWires:
+    """Test the methods and attributes of the FreeWires class"""
+
+    def test_init(self):
+        """Test that the FreeWires class is instantiated as expected when there is no recording."""
+
+        for i in range(3):
+            assert FreeWires(i).num_wires == i
+
+    def test_repr(self):
+        """Test that correct representation is returned for FreeWires class"""
+
+        wires = 3
+        exp_string = f"FreeWires({wires})"
+        assert repr(FreeWires(wires)) == exp_string
+
+    def test_init_recording(self):
+        """Test that the FreeWires class is instantiated as expected when there is active recording."""
+        with qml.queuing.AnnotatedQueue() as q:
+            ops = [FreeWires(2), FreeWires(4), FreeWires(8)]
+
+        assert q.queue == ops