From 20eef13128e31289fcb60e12a79bed68740f4ef7 Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Tue, 30 Apr 2024 10:51:41 -0400
Subject: [PATCH 01/24] implement breakpoint()

---
 pennylane/__init__.py       |  2 +-
 pennylane/debugging.py      | 52 +++++++++++++++++++++++++++++++++++++
 pennylane/workflow/qnode.py | 10 +++++++
 3 files changed, 63 insertions(+), 1 deletion(-)

diff --git a/pennylane/__init__.py b/pennylane/__init__.py
index 07aa57f556d..3aa521b31df 100644
--- a/pennylane/__init__.py
+++ b/pennylane/__init__.py
@@ -122,7 +122,7 @@
 )
 from pennylane.ops.identity import I
 from pennylane.optimize import *
-from pennylane.debugging import snapshots
+from pennylane.debugging import snapshots, breakpoint
 from pennylane.shadows import ClassicalShadow
 from pennylane.qcut import cut_circuit, cut_circuit_mc
 import pennylane.pulse
diff --git a/pennylane/debugging.py b/pennylane/debugging.py
index f6ab56c13b6..99e6473accf 100644
--- a/pennylane/debugging.py
+++ b/pennylane/debugging.py
@@ -14,6 +14,9 @@
 """
 This module contains functionality for debugging quantum programs on simulator devices.
 """
+import pdb
+import sys
+
 import pennylane as qml
 from pennylane import DeviceError
 
@@ -105,3 +108,52 @@ def get_snapshots(*args, **kwargs):
         return dbg.snapshots
 
     return get_snapshots
+
+
+class PLDB(pdb.Pdb):
+    """Custom debugging class integrated with Pdb."""
+
+    __active_dev = []
+
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        self.prompt = "[pldb]: "
+
+    @classmethod
+    def valid_context(cls):
+        """Determine if the debugger is called in a valid context.
+
+        Raises:
+            TypeError: Can't call breakpoint outside of a qnode execution
+            TypeError: Device not supported with breakpoint
+        """
+
+        if not qml.queuing.QueuingManager.recording() or not cls.is_active_dev():
+            raise TypeError("Can't call breakpoint outside of a qnode execution")
+
+        if cls.get_active_device().name not in ("default.qubit", "lightning.qubit"):
+            raise TypeError("Device not supported with breakpoint")
+
+    @classmethod
+    def add_device(cls, dev):
+        cls.__active_dev.append(dev)
+
+    @classmethod
+    def get_active_device(cls):
+        return cls.__active_dev[0]
+
+    @classmethod
+    def is_active_dev(cls):
+        return bool(cls.__active_dev)
+
+    @classmethod
+    def reset_active_dev(cls):
+        print("was reset?")
+        cls.__active_dev = []
+
+
+def breakpoint():
+    """Launch the custom PennyLane debugger."""
+    PLDB.valid_context()
+    debugger = PLDB()
+    debugger.set_trace(sys._getframe().f_back)
diff --git a/pennylane/workflow/qnode.py b/pennylane/workflow/qnode.py
index 7941fa89fe3..147c773a199 100644
--- a/pennylane/workflow/qnode.py
+++ b/pennylane/workflow/qnode.py
@@ -27,6 +27,7 @@
 from pennylane import Device
 from pennylane.measurements import CountsMP, MidMeasureMP, Shots
 from pennylane.tape import QuantumTape, QuantumScript
+from pennylane.debugging import PLDB
 
 from .execution import INTERFACE_MAP, SUPPORTED_INTERFACES
 
@@ -915,10 +916,19 @@ def construct(self, args, kwargs):  # pylint: disable=too-many-branches
         if old_interface == "auto":
             self.interface = qml.math.get_interface(*args, *list(kwargs.values()))
 
+        # Before constructing the tape, we pass the device to the
+        # debugger to ensure they are compatible if there are any
+        # breakpoints in the circuit
+        if PLDB.is_active_dev():
+            PLDB.reset_active_dev()
+
+        PLDB.add_device(self.device)
+
         with qml.queuing.AnnotatedQueue() as q:
             self._qfunc_output = self.func(*args, **kwargs)
 
         self._tape = QuantumScript.from_queue(q, shots)
+        PLDB.reset_active_dev()  # reset active device on the debugger after queuing
 
         params = self.tape.get_parameters(trainable_only=False)
         self.tape.trainable_params = qml.math.get_trainable_indices(params)

From c814e7ed0edde5ae97069019665fa5694ee07d62 Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Tue, 30 Apr 2024 13:38:47 -0400
Subject: [PATCH 02/24] added docstring and tests

---
 pennylane/debugging.py  |  29 ++++++++--
 tests/test_debugging.py | 116 ++++++++++++++++++++++++++++++++++++++++
 2 files changed, 142 insertions(+), 3 deletions(-)

diff --git a/pennylane/debugging.py b/pennylane/debugging.py
index 99e6473accf..bdef1432071 100644
--- a/pennylane/debugging.py
+++ b/pennylane/debugging.py
@@ -116,6 +116,7 @@ class PLDB(pdb.Pdb):
     __active_dev = []
 
     def __init__(self, *args, **kwargs):
+        """Initialize the debugger, and set custom prompt string."""
         super().__init__(*args, **kwargs)
         self.prompt = "[pldb]: "
 
@@ -136,24 +137,46 @@ def valid_context(cls):
 
     @classmethod
     def add_device(cls, dev):
+        """Add a device to the global active device list.
+
+        Args:
+            dev (Union[Device, "qml.devices.Device"]): The active device
+        """
         cls.__active_dev.append(dev)
 
     @classmethod
     def get_active_device(cls):
+        """Return the active device.
+
+        Raises:
+            ValueError: No active device to get
+
+        Returns:
+            Union[Device, "qml.devices.Device"]: The active device
+        """
+        if not cls.is_active_dev():
+            raise ValueError("No active device to get")
+
         return cls.__active_dev[0]
 
     @classmethod
     def is_active_dev(cls):
+        """Determine if there is currently an active device.
+
+        Returns:
+            bool: True if there is an active device
+        """
         return bool(cls.__active_dev)
 
     @classmethod
     def reset_active_dev(cls):
-        print("was reset?")
+        """Reset the global active device list (to empty)."""
         cls.__active_dev = []
 
 
 def breakpoint():
     """Launch the custom PennyLane debugger."""
-    PLDB.valid_context()
+    PLDB.valid_context()  # Ensure its being executed in a valid context
+
     debugger = PLDB()
-    debugger.set_trace(sys._getframe().f_back)
+    debugger.set_trace(sys._getframe().f_back)  # pylint: disable=protected-access
diff --git a/tests/test_debugging.py b/tests/test_debugging.py
index e0d368e5679..925d2a6f8e9 100644
--- a/tests/test_debugging.py
+++ b/tests/test_debugging.py
@@ -18,6 +18,8 @@
 import numpy as np
 import pennylane as qml
 
+from pennylane.debugging import PLDB
+
 
 class TestSnapshot:
     """Test the Snapshot instruction for simulators."""
@@ -375,3 +377,117 @@ def circuit():
             return qml.expval(qml.PauliZ(0))
 
         qml.snapshots(circuit)()
+
+
+# pylint: disable=protected-access
+class TestPLDB:
+    """Test the interactive debugging integration"""
+
+    def test_pldb_init(self):
+        """Test that PLDB initializes correctly"""
+        debugger = PLDB()
+        assert debugger.prompt == "[pldb]: "
+        assert getattr(debugger, "_PLDB__active_dev") == []
+
+    def test_valid_context_outside_qnode(self):
+        """Test that valid_context raises an error when breakpoint
+        is called outside of a qnode execution."""
+
+        with pytest.raises(TypeError, match="Can't call breakpoint outside of a qnode execution"):
+            with qml.queuing.AnnotatedQueue() as q:
+                qml.X(0)
+                qml.breakpoint()
+                qml.Hadamard(0)
+
+        def my_qfunc():
+            qml.X(0)
+            qml.breakpoint()
+            qml.Hadamard(0)
+            return qml.expval(qml.Z(0))
+
+        with pytest.raises(TypeError, match="Can't call breakpoint outside of a qnode execution"):
+            _ = my_qfunc()
+
+    def test_valid_context_not_compatible_device(self):
+        """Test that valid_context raises an error when breakpoint
+        is called outside of a qnode execution."""
+        dev = qml.device("default.mixed", wires=2)
+
+        @qml.qnode(dev)
+        def my_circ():
+            qml.X(0)
+            qml.breakpoint()
+            qml.Hadamard(0)
+            return qml.expva(qml.Z(0))
+
+        with pytest.raises(TypeError, match="Device not supported with breakpoint"):
+            _ = my_circ()
+
+        PLDB.reset_active_dev()
+
+    def test_add_device(self):
+        """Test that we can add a device to the global active device list."""
+        assert getattr(PLDB, "_PLDB__active_dev") == []
+
+        dev1, dev2, dev3 = (
+            qml.device("default.qubit", wires=3),
+            qml.device("default.qubit"),
+            qml.device("lightning.qubit", wires=1),
+        )
+
+        PLDB.add_device(dev1)
+        assert getattr(PLDB, "_PLDB__active_dev") == [dev1]
+
+        PLDB.add_device(dev2)
+        PLDB.add_device(dev3)
+        debugger_active_devs = getattr(PLDB, "_PLDB__active_dev")
+
+        for active_dev, d in zip(debugger_active_devs, [dev1, dev2, dev3]):
+            assert active_dev is d
+
+        PLDB.reset_active_dev()  # clean up the debugger active devices
+
+    dev_names = (
+        "default.qubit",
+        "lightning.qubit",
+    )
+
+    @pytest.mark.parametrize("device_name", dev_names)
+    def test_get_active_device(self, device_name):
+        """Test that we can accses the active device."""
+        dev = qml.device(device_name, wires=2)
+        PLDB.add_device(dev)
+
+        debugger_dev = PLDB.get_active_device()
+        assert debugger_dev is dev
+
+        PLDB.reset_active_dev()
+
+    def test_get_active_device_error_when_no_active_device(self):
+        """Test that an error is raised if we try to get
+        the active device when there are no active devices."""
+        assert getattr(PLDB, "_PLDB__active_dev") == []
+
+        with pytest.raises(ValueError, match="No active device to get"):
+            _ = PLDB.get_active_device()
+
+    @pytest.mark.parametrize("device_name", dev_names)
+    def test_reset_active_device(self, device_name):
+        """Test that we can rest the global active device list."""
+        dev = qml.device(device_name, wires=2)
+        PLDB.add_device(dev)
+        assert getattr(PLDB, "_PLDB__active_dev") == [dev]
+
+        PLDB.reset_active_dev()
+        assert getattr(PLDB, "_PLDB__active_dev") == []
+
+    def test_is_active_device(self):
+        """Test that we can determine if there is an active device."""
+        assert getattr(PLDB, "_PLDB__active_dev") == []
+
+        dev = qml.device("default.qubit")
+        PLDB.add_device(dev)
+        assert PLDB.is_active_dev() == True
+
+        PLDB.reset_active_dev()
+        assert PLDB.is_active_dev() == False

From 508bb8be65fc032dcddbfa7ee6f6b90bbc22a8a2 Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Tue, 30 Apr 2024 13:48:41 -0400
Subject: [PATCH 03/24] format

---
 tests/test_debugging.py | 6 +++---
 1 file changed, 3 insertions(+), 3 deletions(-)

diff --git a/tests/test_debugging.py b/tests/test_debugging.py
index 925d2a6f8e9..6b4a9cc62e1 100644
--- a/tests/test_debugging.py
+++ b/tests/test_debugging.py
@@ -394,7 +394,7 @@ def test_valid_context_outside_qnode(self):
         is called outside of a qnode execution."""
 
         with pytest.raises(TypeError, match="Can't call breakpoint outside of a qnode execution"):
-            with qml.queuing.AnnotatedQueue() as q:
+            with qml.queuing.AnnotatedQueue() as _:
                 qml.X(0)
                 qml.breakpoint()
                 qml.Hadamard(0)
@@ -487,7 +487,7 @@ def test_is_active_device(self):
 
         dev = qml.device("default.qubit")
         PLDB.add_device(dev)
-        assert PLDB.is_active_dev() == True
+        assert PLDB.is_active_dev() is True
 
         PLDB.reset_active_dev()
-        assert PLDB.is_active_dev() == False
+        assert PLDB.is_active_dev() is False

From 18bffe2a9fc73ef7d044548acc9f5a89be755cf7 Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Fri, 10 May 2024 17:01:06 -0400
Subject: [PATCH 04/24] Add missing test coverage

---
 pennylane/workflow/qnode.py |  2 +-
 tests/test_debugging.py     | 38 ++++++++++++++++++++++++++++++++++++-
 2 files changed, 38 insertions(+), 2 deletions(-)

diff --git a/pennylane/workflow/qnode.py b/pennylane/workflow/qnode.py
index ac031ab2daa..ab32434c711 100644
--- a/pennylane/workflow/qnode.py
+++ b/pennylane/workflow/qnode.py
@@ -25,9 +25,9 @@
 
 import pennylane as qml
 from pennylane import Device
+from pennylane.debugging import PLDB
 from pennylane.measurements import CountsMP, MidMeasureMP, Shots
 from pennylane.tape import QuantumScript, QuantumTape
-from pennylane.debugging import PLDB
 
 from .execution import INTERFACE_MAP, SUPPORTED_INTERFACES
 
diff --git a/tests/test_debugging.py b/tests/test_debugging.py
index 8b472e8c968..44d2061cf6d 100644
--- a/tests/test_debugging.py
+++ b/tests/test_debugging.py
@@ -14,11 +14,12 @@
 """
 Unit tests for the debugging module.
 """
+from unittest.mock import patch
+
 import numpy as np
 import pytest
 
 import pennylane as qml
-
 from pennylane.debugging import PLDB
 
 
@@ -492,3 +493,38 @@ def test_is_active_device(self):
 
         PLDB.reset_active_dev()
         assert PLDB.is_active_dev() is False
+
+
+@patch.object(PLDB, "set_trace")
+def test_breakpoint_integration(mock_method):
+    """Test that qml.breakpoint behaves as execpted"""
+    dev = qml.device("default.qubit")
+
+    @qml.qnode(dev)
+    def my_circ():
+        qml.Hadamard(0)
+        qml.CNOT([0, 1])
+        qml.breakpoint()
+        return qml.expval(qml.Z(1))
+
+    mock_method.assert_not_called()  # Did not hit breakpoint
+    my_circ()
+    mock_method.assert_called_once()  # Hit breakpoint once.
+
+
+@patch.object(PLDB, "set_trace")
+def test_breakpoint_integration_with_valid_context_error(mock_method):
+    """Test that the PLDB.valid_context() integrates well with qml.breakpoint"""
+    dev = qml.device("default.mixed", wires=2)
+
+    @qml.qnode(dev)
+    def my_circ():
+        qml.Hadamard(0)
+        qml.CNOT([0, 1])
+        qml.breakpoint()
+        return qml.expval(qml.Z(1))
+
+    with pytest.raises(TypeError, match="Device not supported with breakpoint"):
+        _ = my_circ()
+
+    mock_method.assert_not_called()  # Error was raised before we triggered breakpoint

From ac476facd3de59ff44dbebe839744c2984a33019 Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Tue, 28 May 2024 12:10:48 -0400
Subject: [PATCH 05/24] Added support for measurements

---
 pennylane/debugging.py  | 68 +++++++++++++++++++++++++++++++++++++++++
 tests/test_debugging.py |  2 +-
 2 files changed, 69 insertions(+), 1 deletion(-)

diff --git a/pennylane/debugging.py b/pennylane/debugging.py
index 10549d82849..e41effba8b5 100644
--- a/pennylane/debugging.py
+++ b/pennylane/debugging.py
@@ -16,6 +16,7 @@
 """
 import pdb
 import sys
+import copy
 
 import pennylane as qml
 from pennylane import DeviceError
@@ -173,6 +174,16 @@ def reset_active_dev(cls):
         """Reset the global active device list (to empty)."""
         cls.__active_dev = []
 
+    @classmethod
+    def _execute(cls, batch_tapes):
+        """Execute tape on the active device"""
+        dev = cls.get_active_device()
+
+        program, new_config = dev.preprocess()
+        new_batch, fn = program(batch_tapes)
+
+        return fn(dev.execute(new_batch, new_config))
+
 
 def breakpoint():
     """Launch the custom PennyLane debugger."""
@@ -180,3 +191,60 @@ def breakpoint():
 
     debugger = PLDB()
     debugger.set_trace(sys._getframe().f_back)  # pylint: disable=protected-access
+
+
+def state():
+    """Compute the state of the quantum circuit.
+
+    Returns:
+        Array(complex): Quantum state of the circuit.
+    """
+    with qml.queuing.QueuingManager.stop_recording():
+        m = qml.state()
+
+    return _measure(m)
+
+
+def expval(op):
+    """Compute the expectation value of an observable.
+
+    Args:
+        op (Operator): The observable to compute the expectation value for.
+
+    Returns:
+        complex: Quantum state of the circuit.
+    """
+
+    qml.queuing.QueuingManager.active_context().remove(op)  # ensure we didn't accidentally queue op
+
+    with qml.queuing.QueuingManager.stop_recording():
+        m = qml.expval(op)
+
+    return _measure(m)
+
+
+def _measure(measurement):
+    """Perform the measurement.
+
+    Args:
+        measurement (MeasurementProcess): The type of measurement to be performed
+
+    Returns:
+        tuple(complex): Results from the measurement
+    """
+    active_queue = qml.queuing.QueuingManager.active_context()
+    copied_queue = copy.deepcopy(active_queue)
+    
+    copied_queue.append(measurement)
+    tape = qml.tape.QuantumScript.from_queue(copied_queue)
+    return PLDB._execute((tape,))
+  
+
+def tape():
+    """Access the tape of the quantum circuit.
+
+    Returns:
+        QuantumScript: The quantum script representing the circuit.
+    """
+    active_queue = qml.queuing.QueuingManager.active_context()
+    return qml.tape.QuantumScript.from_queue(active_queue)
diff --git a/tests/test_debugging.py b/tests/test_debugging.py
index 44d2061cf6d..2a19db4ef4f 100644
--- a/tests/test_debugging.py
+++ b/tests/test_debugging.py
@@ -412,7 +412,7 @@ def my_qfunc():
 
     def test_valid_context_not_compatible_device(self):
         """Test that valid_context raises an error when breakpoint
-        is called outside of a qnode execution."""
+        is called in a qnode with an incompatible device."""
         dev = qml.device("default.mixed", wires=2)
 
         @qml.qnode(dev)

From 5b1ad2c63367fc91ed3ff54f68a2e7ae2647acc8 Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Tue, 28 May 2024 17:22:07 -0400
Subject: [PATCH 06/24] Added tests

---
 pennylane/debugging.py  |  6 +--
 tests/test_debugging.py | 89 +++++++++++++++++++++++++++++++++++++++++
 2 files changed, 92 insertions(+), 3 deletions(-)

diff --git a/pennylane/debugging.py b/pennylane/debugging.py
index e41effba8b5..c8d8db37fed 100644
--- a/pennylane/debugging.py
+++ b/pennylane/debugging.py
@@ -14,9 +14,9 @@
 """
 This module contains functionality for debugging quantum programs on simulator devices.
 """
+import copy
 import pdb
 import sys
-import copy
 
 import pennylane as qml
 from pennylane import DeviceError
@@ -234,11 +234,11 @@ def _measure(measurement):
     """
     active_queue = qml.queuing.QueuingManager.active_context()
     copied_queue = copy.deepcopy(active_queue)
-    
+
     copied_queue.append(measurement)
     tape = qml.tape.QuantumScript.from_queue(copied_queue)
     return PLDB._execute((tape,))
-  
+
 
 def tape():
     """Access the tape of the quantum circuit.
diff --git a/tests/test_debugging.py b/tests/test_debugging.py
index 2a19db4ef4f..7d317362bc7 100644
--- a/tests/test_debugging.py
+++ b/tests/test_debugging.py
@@ -20,6 +20,7 @@
 import pytest
 
 import pennylane as qml
+from pennylane import numpy as qnp
 from pennylane.debugging import PLDB
 
 
@@ -494,6 +495,94 @@ def test_is_active_device(self):
         PLDB.reset_active_dev()
         assert PLDB.is_active_dev() is False
 
+    tapes = (
+        qml.tape.QuantumScript(
+            ops=[qml.Hadamard(0), qml.CNOT([0, 1])],
+            measurements=[qml.state()],
+        ),
+        qml.tape.QuantumScript(
+            ops=[qml.Hadamard(0), qml.X(1)],
+            measurements=[qml.expval(qml.Z(1))],
+        ),
+    )
+
+    results = (qnp.array([1 / qnp.sqrt(2), 0, 0, 1 / qnp.sqrt(2)], dtype=complex), qnp.array(-1))
+
+    @pytest.mark.parametrize("tape, expected_result", zip(tapes, results))
+    @pytest.mark.parametrize(
+        "dev", (qml.device("default.qubit", wires=2), qml.device("lightning.qubit", wires=2))
+    )
+    def test_execute(self, dev, tape, expected_result):
+        """Test that the _execute method works as expected."""
+        PLDB.add_device(dev)
+        executed_results = PLDB._execute((tape,))
+        assert qnp.allclose(expected_result, executed_results)
+
+
+def test_tape():
+    """Test that we can access the tape from the active queue."""
+    with qml.queuing.AnnotatedQueue() as queue:
+        qml.X(0)
+        [qml.Hadamard(i) for i in range(3)]
+        qml.Y(1)
+        qml.Z(0)
+        qml.expval(qml.Z(0))
+
+        executed_tape = qml.debugging.tape()
+
+    expected_tape = qml.tape.QuantumScript.from_queue(queue)
+    assert qml.equal(expected_tape, executed_tape)
+
+
+@pytest.mark.parametrize("measurement_process", (qml.expval(qml.Z(0)), qml.state()))
+@patch.object(PLDB, "_execute")
+def test_measure(mock_method, measurement_process):
+    """Test that the private measure function doesn't modify the active queue"""
+    with qml.queuing.AnnotatedQueue() as queue:
+        ops = [qml.X(0), qml.Y(1), qml.Z(0)] + [qml.Hadamard(i) for i in range(3)]
+        measurements = [qml.expval(qml.X(2)), qml.state(), qml.probs(), qml.var(qml.Z(3))]
+        _ = qml.debugging._measure(measurement_process)
+
+    executed_tape = qml.tape.QuantumScript.from_queue(queue)
+    expected_tape = qml.tape.QuantumScript(ops, measurements)
+
+    assert qml.equal(expected_tape, executed_tape)  # no unexpected queuing
+
+    expected_debugging_tape = qml.tape.QuantumScript(ops, measurements + [measurement_process])
+    executed_debugging_tape = mock_method.call_args.args[0][0]
+
+    assert qml.equal(
+        expected_debugging_tape, executed_debugging_tape
+    )  # _execute was called with new measurements
+
+
+@patch.object(PLDB, "_execute")
+def test_state(_mock_method):
+    """Test that the state function works as expected."""
+    with qml.queuing.AnnotatedQueue() as queue:
+        qml.RX(1.23, 0)
+        qml.RY(0.45, 2)
+        qml.probs()
+
+        _ = qml.debugging.state()
+
+    assert qml.state() not in queue
+
+
+@patch.object(PLDB, "_execute")
+def test_expval(_mock_method):
+    """Test that the expval function works as expected."""
+    for op in [qml.X(0), qml.Y(1), qml.Z(2), qml.Hadamard(0)]:
+        with qml.queuing.AnnotatedQueue() as queue:
+            qml.RX(1.23, 0)
+            qml.RY(0.45, 2)
+            qml.probs()
+
+            _ = qml.debugging.expval(op)
+
+        assert op not in queue
+        assert qml.expval(op) not in queue
+
 
 @patch.object(PLDB, "set_trace")
 def test_breakpoint_integration(mock_method):

From f7feadbe3be2b8fecd7c78c237c40161196d4616 Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Thu, 30 May 2024 09:56:36 -0400
Subject: [PATCH 07/24] Address code review comments

---
 pennylane/debugging.py  | 12 ++++++------
 tests/test_debugging.py | 16 ++++++++++------
 2 files changed, 16 insertions(+), 12 deletions(-)

diff --git a/pennylane/debugging.py b/pennylane/debugging.py
index 10549d82849..f732a3711ee 100644
--- a/pennylane/debugging.py
+++ b/pennylane/debugging.py
@@ -125,15 +125,15 @@ def valid_context(cls):
         """Determine if the debugger is called in a valid context.
 
         Raises:
-            TypeError: Can't call breakpoint outside of a qnode execution
-            TypeError: Device not supported with breakpoint
+            RuntimeError: Can't call breakpoint outside of a qnode execution
+            TypeError: Breakpoints not supported on this device
         """
 
         if not qml.queuing.QueuingManager.recording() or not cls.is_active_dev():
-            raise TypeError("Can't call breakpoint outside of a qnode execution")
+            raise RuntimeError("Can't call breakpoint outside of a qnode execution")
 
         if cls.get_active_device().name not in ("default.qubit", "lightning.qubit"):
-            raise TypeError("Device not supported with breakpoint")
+            raise TypeError("Breakpoints not supported on this device")
 
     @classmethod
     def add_device(cls, dev):
@@ -149,13 +149,13 @@ def get_active_device(cls):
         """Return the active device.
 
         Raises:
-            ValueError: No active device to get
+            RuntimeError: No active device to get
 
         Returns:
             Union[Device, "qml.devices.Device"]: The active device
         """
         if not cls.is_active_dev():
-            raise ValueError("No active device to get")
+            raise RuntimeError("No active device to get")
 
         return cls.__active_dev[0]
 
diff --git a/tests/test_debugging.py b/tests/test_debugging.py
index 44d2061cf6d..e656dcf664b 100644
--- a/tests/test_debugging.py
+++ b/tests/test_debugging.py
@@ -395,7 +395,9 @@ def test_valid_context_outside_qnode(self):
         """Test that valid_context raises an error when breakpoint
         is called outside of a qnode execution."""
 
-        with pytest.raises(TypeError, match="Can't call breakpoint outside of a qnode execution"):
+        with pytest.raises(
+            RuntimeError, match="Can't call breakpoint outside of a qnode execution"
+        ):
             with qml.queuing.AnnotatedQueue() as _:
                 qml.X(0)
                 qml.breakpoint()
@@ -407,12 +409,14 @@ def my_qfunc():
             qml.Hadamard(0)
             return qml.expval(qml.Z(0))
 
-        with pytest.raises(TypeError, match="Can't call breakpoint outside of a qnode execution"):
+        with pytest.raises(
+            RuntimeError, match="Can't call breakpoint outside of a qnode execution"
+        ):
             _ = my_qfunc()
 
     def test_valid_context_not_compatible_device(self):
         """Test that valid_context raises an error when breakpoint
-        is called outside of a qnode execution."""
+        is called with an un-supported device."""
         dev = qml.device("default.mixed", wires=2)
 
         @qml.qnode(dev)
@@ -422,7 +426,7 @@ def my_circ():
             qml.Hadamard(0)
             return qml.expva(qml.Z(0))
 
-        with pytest.raises(TypeError, match="Device not supported with breakpoint"):
+        with pytest.raises(TypeError, match="Breakpoints not supported on this device"):
             _ = my_circ()
 
         PLDB.reset_active_dev()
@@ -470,7 +474,7 @@ def test_get_active_device_error_when_no_active_device(self):
         the active device when there are no active devices."""
         assert getattr(PLDB, "_PLDB__active_dev") == []
 
-        with pytest.raises(ValueError, match="No active device to get"):
+        with pytest.raises(RuntimeError, match="No active device to get"):
             _ = PLDB.get_active_device()
 
     @pytest.mark.parametrize("device_name", dev_names)
@@ -524,7 +528,7 @@ def my_circ():
         qml.breakpoint()
         return qml.expval(qml.Z(1))
 
-    with pytest.raises(TypeError, match="Device not supported with breakpoint"):
+    with pytest.raises(TypeError, match="Breakpoints not supported on this device"):
         _ = my_circ()
 
     mock_method.assert_not_called()  # Error was raised before we triggered breakpoint

From 2e00723e61752d23fa0973296d44a160576941c6 Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Thu, 30 May 2024 16:35:09 -0400
Subject: [PATCH 08/24] add context manager + address code review comments

---
 pennylane/debugging.py      | 28 ++++++++++++++++++++++------
 pennylane/workflow/qnode.py | 13 ++++---------
 tests/test_debugging.py     | 36 +++++++++++++++++++++++-------------
 3 files changed, 49 insertions(+), 28 deletions(-)

diff --git a/pennylane/debugging.py b/pennylane/debugging.py
index f732a3711ee..9aaae8d1b95 100644
--- a/pennylane/debugging.py
+++ b/pennylane/debugging.py
@@ -16,6 +16,7 @@
 """
 import pdb
 import sys
+from contextlib import contextmanager
 
 import pennylane as qml
 from pennylane import DeviceError
@@ -113,7 +114,7 @@ def get_snapshots(*args, **kwargs):
 class PLDB(pdb.Pdb):
     """Custom debugging class integrated with Pdb."""
 
-    __active_dev = []
+    __active_dev = None
 
     def __init__(self, *args, **kwargs):
         """Initialize the debugger, and set custom prompt string."""
@@ -137,12 +138,12 @@ def valid_context(cls):
 
     @classmethod
     def add_device(cls, dev):
-        """Add a device to the global active device list.
+        """Update the global active device.
 
         Args:
             dev (Union[Device, "qml.devices.Device"]): The active device
         """
-        cls.__active_dev.append(dev)
+        cls.__active_dev = dev
 
     @classmethod
     def get_active_device(cls):
@@ -157,7 +158,7 @@ def get_active_device(cls):
         if not cls.is_active_dev():
             raise RuntimeError("No active device to get")
 
-        return cls.__active_dev[0]
+        return cls.__active_dev
 
     @classmethod
     def is_active_dev(cls):
@@ -171,12 +172,27 @@ def is_active_dev(cls):
     @classmethod
     def reset_active_dev(cls):
         """Reset the global active device list (to empty)."""
-        cls.__active_dev = []
+        cls.__active_dev = None
+
+
+@contextmanager
+def pldb_device_manager(device):
+    """Context manager to automatically set and reset active
+    device on the PLDB debugger.
+
+    Args:
+        device (Union[Device, "qml.devices.Device"]): The active device instance
+    """
+    try:
+        PLDB.add_device(device)
+        yield
+    finally:
+        PLDB.reset_active_dev()
 
 
 def breakpoint():
     """Launch the custom PennyLane debugger."""
     PLDB.valid_context()  # Ensure its being executed in a valid context
 
-    debugger = PLDB()
+    debugger = PLDB(skip=["pennylane.*"])  # skip internals when stepping through trace
     debugger.set_trace(sys._getframe().f_back)  # pylint: disable=protected-access
diff --git a/pennylane/workflow/qnode.py b/pennylane/workflow/qnode.py
index ab32434c711..8b2c647e1bd 100644
--- a/pennylane/workflow/qnode.py
+++ b/pennylane/workflow/qnode.py
@@ -25,7 +25,7 @@
 
 import pennylane as qml
 from pennylane import Device
-from pennylane.debugging import PLDB
+from pennylane.debugging import pldb_device_manager
 from pennylane.measurements import CountsMP, MidMeasureMP, Shots
 from pennylane.tape import QuantumScript, QuantumTape
 
@@ -929,16 +929,11 @@ def construct(self, args, kwargs):  # pylint: disable=too-many-branches
         # Before constructing the tape, we pass the device to the
         # debugger to ensure they are compatible if there are any
         # breakpoints in the circuit
-        if PLDB.is_active_dev():
-            PLDB.reset_active_dev()
-
-        PLDB.add_device(self.device)
-
-        with qml.queuing.AnnotatedQueue() as q:
-            self._qfunc_output = self.func(*args, **kwargs)
+        with pldb_device_manager(self.device) as _:
+            with qml.queuing.AnnotatedQueue() as q:
+                self._qfunc_output = self.func(*args, **kwargs)
 
         self._tape = QuantumScript.from_queue(q, shots)
-        PLDB.reset_active_dev()  # reset active device on the debugger after queuing
 
         params = self.tape.get_parameters(trainable_only=False)
         self.tape.trainable_params = qml.math.get_trainable_indices(params)
diff --git a/tests/test_debugging.py b/tests/test_debugging.py
index e656dcf664b..a91503573d3 100644
--- a/tests/test_debugging.py
+++ b/tests/test_debugging.py
@@ -20,7 +20,7 @@
 import pytest
 
 import pennylane as qml
-from pennylane.debugging import PLDB
+from pennylane.debugging import PLDB, pldb_device_manager
 
 
 class TestSnapshot:
@@ -389,7 +389,7 @@ def test_pldb_init(self):
         """Test that PLDB initializes correctly"""
         debugger = PLDB()
         assert debugger.prompt == "[pldb]: "
-        assert getattr(debugger, "_PLDB__active_dev") == []
+        assert getattr(debugger, "_PLDB__active_dev") is None
 
     def test_valid_context_outside_qnode(self):
         """Test that valid_context raises an error when breakpoint
@@ -433,7 +433,7 @@ def my_circ():
 
     def test_add_device(self):
         """Test that we can add a device to the global active device list."""
-        assert getattr(PLDB, "_PLDB__active_dev") == []
+        assert getattr(PLDB, "_PLDB__active_dev") is None
 
         dev1, dev2, dev3 = (
             qml.device("default.qubit", wires=3),
@@ -442,14 +442,12 @@ def test_add_device(self):
         )
 
         PLDB.add_device(dev1)
-        assert getattr(PLDB, "_PLDB__active_dev") == [dev1]
+        assert getattr(PLDB, "_PLDB__active_dev") == dev1
 
-        PLDB.add_device(dev2)
-        PLDB.add_device(dev3)
-        debugger_active_devs = getattr(PLDB, "_PLDB__active_dev")
+        PLDB.add_device(dev2)  # overwrites dev1
+        PLDB.add_device(dev3)  # overwrites dev2
 
-        for active_dev, d in zip(debugger_active_devs, [dev1, dev2, dev3]):
-            assert active_dev is d
+        assert getattr(PLDB, "_PLDB__active_dev") == dev3
 
         PLDB.reset_active_dev()  # clean up the debugger active devices
 
@@ -472,7 +470,7 @@ def test_get_active_device(self, device_name):
     def test_get_active_device_error_when_no_active_device(self):
         """Test that an error is raised if we try to get
         the active device when there are no active devices."""
-        assert getattr(PLDB, "_PLDB__active_dev") == []
+        assert getattr(PLDB, "_PLDB__active_dev") == None
 
         with pytest.raises(RuntimeError, match="No active device to get"):
             _ = PLDB.get_active_device()
@@ -482,14 +480,14 @@ def test_reset_active_device(self, device_name):
         """Test that we can rest the global active device list."""
         dev = qml.device(device_name, wires=2)
         PLDB.add_device(dev)
-        assert getattr(PLDB, "_PLDB__active_dev") == [dev]
+        assert getattr(PLDB, "_PLDB__active_dev") == dev
 
         PLDB.reset_active_dev()
-        assert getattr(PLDB, "_PLDB__active_dev") == []
+        assert getattr(PLDB, "_PLDB__active_dev") == None
 
     def test_is_active_device(self):
         """Test that we can determine if there is an active device."""
-        assert getattr(PLDB, "_PLDB__active_dev") == []
+        assert getattr(PLDB, "_PLDB__active_dev") == None
 
         dev = qml.device("default.qubit")
         PLDB.add_device(dev)
@@ -499,6 +497,18 @@ def test_is_active_device(self):
         assert PLDB.is_active_dev() is False
 
 
+@pytest.mark.parametrize("device_name", ("default.qubit", "lightning.qubit"))
+def test_pldb_device_manager(device_name):
+    """Test that the context manager works as expected."""
+    assert getattr(PLDB, "_PLDB__active_dev") == None
+    dev = qml.device(device_name, wires=2)
+
+    with pldb_device_manager(dev) as _:
+        assert getattr(PLDB, "_PLDB__active_dev") == dev
+
+    assert getattr(PLDB, "_PLDB__active_dev") == None
+
+
 @patch.object(PLDB, "set_trace")
 def test_breakpoint_integration(mock_method):
     """Test that qml.breakpoint behaves as execpted"""

From a897a476c549aafae33a2ddfdeff360d7903bb11 Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Thu, 30 May 2024 16:40:31 -0400
Subject: [PATCH 09/24] [skip ci]


From 6751e8779d5bc6d0e82522dfa0c1a61f2270b399 Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Fri, 31 May 2024 10:06:30 -0400
Subject: [PATCH 10/24] fix typo

---
 pennylane/workflow/qnode.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/pennylane/workflow/qnode.py b/pennylane/workflow/qnode.py
index 9c0135cea45..a19fb8f3aae 100644
--- a/pennylane/workflow/qnode.py
+++ b/pennylane/workflow/qnode.py
@@ -25,8 +25,8 @@
 
 import pennylane as qml
 from pennylane import Device
-from pennylane.logging import debug_logger
 from pennylane.debugging import pldb_device_manager
+from pennylane.logging import debug_logger
 from pennylane.measurements import CountsMP, MidMeasureMP, Shots
 from pennylane.tape import QuantumScript, QuantumTape
 

From f537c39b2fd0e41ea29408d4b8048f25033bfa74 Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Fri, 31 May 2024 10:08:06 -0400
Subject: [PATCH 11/24] pull class method back into class

---
 pennylane/debugging.py | 20 ++++++++++----------
 1 file changed, 10 insertions(+), 10 deletions(-)

diff --git a/pennylane/debugging.py b/pennylane/debugging.py
index b6186c174e5..316ecccd2c9 100644
--- a/pennylane/debugging.py
+++ b/pennylane/debugging.py
@@ -175,6 +175,16 @@ def reset_active_dev(cls):
         """Reset the global active device list (to empty)."""
         cls.__active_dev = None
 
+    @classmethod
+    def _execute(cls, batch_tapes):
+        """Execute tape on the active device"""
+        dev = cls.get_active_device()
+
+        program, new_config = dev.preprocess()
+        new_batch, fn = program(batch_tapes)
+
+        return fn(dev.execute(new_batch, new_config))
+
 
 @contextmanager
 def pldb_device_manager(device):
@@ -190,16 +200,6 @@ def pldb_device_manager(device):
     finally:
         PLDB.reset_active_dev()
 
-    @classmethod
-    def _execute(cls, batch_tapes):
-        """Execute tape on the active device"""
-        dev = cls.get_active_device()
-
-        program, new_config = dev.preprocess()
-        new_batch, fn = program(batch_tapes)
-
-        return fn(dev.execute(new_batch, new_config))
-
 
 def breakpoint():
     """Launch the custom PennyLane debugger."""

From 74a2292c80f9991df30ba46c1fd0ed66e5604dd5 Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Mon, 3 Jun 2024 10:39:55 -0400
Subject: [PATCH 12/24] Add probs support

---
 pennylane/debugging.py  | 22 ++++++++++++++++++
 tests/test_debugging.py | 51 ++++++++++++++++++++++++++++++++++++++---
 2 files changed, 70 insertions(+), 3 deletions(-)

diff --git a/pennylane/debugging.py b/pennylane/debugging.py
index 316ecccd2c9..333e09ad131 100644
--- a/pennylane/debugging.py
+++ b/pennylane/debugging.py
@@ -239,6 +239,28 @@ def expval(op):
     return _measure(m)
 
 
+def probs(wires=None, op=None):
+    """Compute the probability distribution for the state.
+    Args:
+        wires (Union[Iterable, int, str, list]): the wires the operation acts on
+        op (Union[Observable, MeasurementValue]): Observable (with a ``diagonalizing_gates``
+            attribute) that rotates the computational basis, or a  ``MeasurementValue``
+            corresponding to mid-circuit measurements.
+
+    Returns:
+        Array(float): The probability distribution of the bitstrings for the wires.
+    """
+    if op:
+        qml.queuing.QueuingManager.active_context().remove(
+            op
+        )  # ensure we didn't accidentally queue op
+
+    with qml.queuing.QueuingManager.stop_recording():
+        m = qml.probs(wires, op)
+
+    return _measure(m)
+
+
 def _measure(measurement):
     """Perform the measurement.
 
diff --git a/tests/test_debugging.py b/tests/test_debugging.py
index 91ea388236e..c59d3d5325c 100644
--- a/tests/test_debugging.py
+++ b/tests/test_debugging.py
@@ -506,9 +506,22 @@ def test_is_active_device(self):
             ops=[qml.Hadamard(0), qml.X(1)],
             measurements=[qml.expval(qml.Z(1))],
         ),
+        qml.tape.QuantumScript(
+            ops=[qml.Hadamard(0), qml.CNOT([0, 1])],
+            measurements=[qml.probs()],
+        ),
+        qml.tape.QuantumScript(
+            ops=[qml.Hadamard(0), qml.CNOT([0, 1])],
+            measurements=[qml.probs(wires=[0])],
+        ),
     )
 
-    results = (qnp.array([1 / qnp.sqrt(2), 0, 0, 1 / qnp.sqrt(2)], dtype=complex), qnp.array(-1))
+    results = (
+        qnp.array([1 / qnp.sqrt(2), 0, 0, 1 / qnp.sqrt(2)], dtype=complex),
+        qnp.array(-1),
+        qnp.array([1 / 2, 0, 0, 1 / 2]),
+        qnp.array([1 / 2, 1 / 2]),
+    )
 
     @pytest.mark.parametrize("tape, expected_result", zip(tapes, results))
     @pytest.mark.parametrize(
@@ -519,6 +532,7 @@ def test_execute(self, dev, tape, expected_result):
         PLDB.add_device(dev)
         executed_results = PLDB._execute((tape,))
         assert qnp.allclose(expected_result, executed_results)
+        PLDB.reset_active_dev()
 
 
 def test_tape():
@@ -564,7 +578,7 @@ def test_state(_mock_method):
     with qml.queuing.AnnotatedQueue() as queue:
         qml.RX(1.23, 0)
         qml.RY(0.45, 2)
-        qml.probs()
+        qml.sample()
 
         _ = qml.debugging.state()
 
@@ -578,7 +592,7 @@ def test_expval(_mock_method):
         with qml.queuing.AnnotatedQueue() as queue:
             qml.RX(1.23, 0)
             qml.RY(0.45, 2)
-            qml.probs()
+            qml.sample()
 
             _ = qml.debugging.expval(op)
 
@@ -586,6 +600,37 @@ def test_expval(_mock_method):
         assert qml.expval(op) not in queue
 
 
+@patch.object(PLDB, "_execute")
+def test_probs_with_op(_mock_method):
+    """Test that the probs function works as expected."""
+
+    for op in [None, qml.X(0), qml.Y(1), qml.Z(2)]:
+        with qml.queuing.AnnotatedQueue() as queue:
+            qml.RX(1.23, 0)
+            qml.RY(0.45, 2)
+            qml.sample()
+
+            _ = qml.debugging.probs(op=op)
+
+        assert op not in queue
+        assert qml.probs(op=op) not in queue
+
+
+@patch.object(PLDB, "_execute")
+def test_probs_with_wires(_mock_method):
+    """Test that the probs function works as expected."""
+
+    for wires in [None, [0, 1], [2]]:
+        with qml.queuing.AnnotatedQueue() as queue:
+            qml.RX(1.23, 0)
+            qml.RY(0.45, 2)
+            qml.sample()
+
+            _ = qml.debugging.probs(wires=wires)
+
+        assert qml.probs(wires=wires) not in queue
+
+
 @pytest.mark.parametrize("device_name", ("default.qubit", "lightning.qubit"))
 def test_pldb_device_manager(device_name):
     """Test that the context manager works as expected."""

From eb8e852d920610f6e03e5b2ccf672ff1578f6d12 Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Mon, 3 Jun 2024 14:58:17 -0400
Subject: [PATCH 13/24] fix output, [skip ci]

---
 pennylane/debugging.py | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/pennylane/debugging.py b/pennylane/debugging.py
index 333e09ad131..6325a810eb6 100644
--- a/pennylane/debugging.py
+++ b/pennylane/debugging.py
@@ -183,7 +183,8 @@ def _execute(cls, batch_tapes):
         program, new_config = dev.preprocess()
         new_batch, fn = program(batch_tapes)
 
-        return fn(dev.execute(new_batch, new_config))
+        # TODO: remove [0] index once compatible with transforms
+        return fn(dev.execute(new_batch, new_config))[0]
 
 
 @contextmanager

From 3f9585da2cb43ea4059051cc4379f284fe38767e Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Tue, 4 Jun 2024 15:52:14 -0400
Subject: [PATCH 14/24] expand state to device wires

---
 pennylane/debugging.py  | 8 +++++++-
 tests/test_debugging.py | 5 +++++
 2 files changed, 12 insertions(+), 1 deletion(-)

diff --git a/pennylane/debugging.py b/pennylane/debugging.py
index 6325a810eb6..49b50f589fc 100644
--- a/pennylane/debugging.py
+++ b/pennylane/debugging.py
@@ -180,8 +180,14 @@ def _execute(cls, batch_tapes):
         """Execute tape on the active device"""
         dev = cls.get_active_device()
 
+        valid_batch, _ = (
+            (batch_tapes, None)
+            if not dev.wires
+            else qml.devices.preprocess.validate_device_wires(batch_tapes, wires=dev.wires)
+        )
+
         program, new_config = dev.preprocess()
-        new_batch, fn = program(batch_tapes)
+        new_batch, fn = program(valid_batch)
 
         # TODO: remove [0] index once compatible with transforms
         return fn(dev.execute(new_batch, new_config))[0]
diff --git a/tests/test_debugging.py b/tests/test_debugging.py
index c59d3d5325c..4b1b8233b6d 100644
--- a/tests/test_debugging.py
+++ b/tests/test_debugging.py
@@ -514,6 +514,10 @@ def test_is_active_device(self):
             ops=[qml.Hadamard(0), qml.CNOT([0, 1])],
             measurements=[qml.probs(wires=[0])],
         ),
+        qml.tape.QuantumScript(
+            ops=[qml.Hadamard(0)],
+            measurements=[qml.state()],
+        ),  # Test that state expands to number of device wires
     )
 
     results = (
@@ -521,6 +525,7 @@ def test_is_active_device(self):
         qnp.array(-1),
         qnp.array([1 / 2, 0, 0, 1 / 2]),
         qnp.array([1 / 2, 1 / 2]),
+        qnp.array([1 / qnp.sqrt(2), 0, 1 / qnp.sqrt(2), 0], dtype=complex),
     )
 
     @pytest.mark.parametrize("tape, expected_result", zip(tapes, results))

From 13f55f00cbd3c5b8723f6606ecc34d1d6e14c4c0 Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Wed, 5 Jun 2024 12:03:18 -0400
Subject: [PATCH 15/24] Apply suggestions from code review

Co-authored-by: Mikhail Andrenkov <mikhail@xanadu.ai>
---
 tests/test_debugging.py | 4 ++--
 1 file changed, 2 insertions(+), 2 deletions(-)

diff --git a/tests/test_debugging.py b/tests/test_debugging.py
index a91503573d3..53110a147ca 100644
--- a/tests/test_debugging.py
+++ b/tests/test_debugging.py
@@ -458,7 +458,7 @@ def test_add_device(self):
 
     @pytest.mark.parametrize("device_name", dev_names)
     def test_get_active_device(self, device_name):
-        """Test that we can accses the active device."""
+        """Test that we can access the active device."""
         dev = qml.device(device_name, wires=2)
         PLDB.add_device(dev)
 
@@ -511,7 +511,7 @@ def test_pldb_device_manager(device_name):
 
 @patch.object(PLDB, "set_trace")
 def test_breakpoint_integration(mock_method):
-    """Test that qml.breakpoint behaves as execpted"""
+    """Test that qml.breakpoint behaves as expected"""
     dev = qml.device("default.qubit")
 
     @qml.qnode(dev)

From 7b3182780a308e94dad1be8d5c224891c356077f Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Wed, 5 Jun 2024 15:22:00 -0400
Subject: [PATCH 16/24] Addressed code review comments

---
 pennylane/debugging.py      | 23 +++++++++++++----------
 pennylane/workflow/qnode.py |  2 +-
 tests/test_debugging.py     | 32 ++++++++++++++------------------
 3 files changed, 28 insertions(+), 29 deletions(-)

diff --git a/pennylane/debugging.py b/pennylane/debugging.py
index 9aaae8d1b95..136847a97e8 100644
--- a/pennylane/debugging.py
+++ b/pennylane/debugging.py
@@ -114,23 +114,26 @@ def get_snapshots(*args, **kwargs):
 class PLDB(pdb.Pdb):
     """Custom debugging class integrated with Pdb."""
 
-    __active_dev = None
+    __active_dev: qml.devices.Device | None = None
 
     def __init__(self, *args, **kwargs):
         """Initialize the debugger, and set custom prompt string."""
         super().__init__(*args, **kwargs)
-        self.prompt = "[pldb]: "
+
+    @property
+    def prompt(self):
+        return "[pldb]: "
 
     @classmethod
     def valid_context(cls):
         """Determine if the debugger is called in a valid context.
 
         Raises:
-            RuntimeError: Can't call breakpoint outside of a qnode execution
-            TypeError: Breakpoints not supported on this device
+            RuntimeError: breakpoint is called outside of a qnode execution
+            TypeError: breakpoints not supported on this device
         """
 
-        if not qml.queuing.QueuingManager.recording() or not cls.is_active_dev():
+        if not qml.queuing.QueuingManager.recording() or not cls.has_active_dev():
             raise RuntimeError("Can't call breakpoint outside of a qnode execution")
 
         if cls.get_active_device().name not in ("default.qubit", "lightning.qubit"):
@@ -141,7 +144,7 @@ def add_device(cls, dev):
         """Update the global active device.
 
         Args:
-            dev (Union[Device, "qml.devices.Device"]): The active device
+            dev (Union[Device, "qml.devices.Device"]): the active device
         """
         cls.__active_dev = dev
 
@@ -155,13 +158,13 @@ def get_active_device(cls):
         Returns:
             Union[Device, "qml.devices.Device"]: The active device
         """
-        if not cls.is_active_dev():
+        if not cls.has_active_dev():
             raise RuntimeError("No active device to get")
 
         return cls.__active_dev
 
     @classmethod
-    def is_active_dev(cls):
+    def has_active_dev(cls):
         """Determine if there is currently an active device.
 
         Returns:
@@ -171,7 +174,7 @@ def is_active_dev(cls):
 
     @classmethod
     def reset_active_dev(cls):
-        """Reset the global active device list (to empty)."""
+        """Reset the global active device variable to None."""
         cls.__active_dev = None
 
 
@@ -181,7 +184,7 @@ def pldb_device_manager(device):
     device on the PLDB debugger.
 
     Args:
-        device (Union[Device, "qml.devices.Device"]): The active device instance
+        device (Union[Device, "qml.devices.Device"]): the active device instance
     """
     try:
         PLDB.add_device(device)
diff --git a/pennylane/workflow/qnode.py b/pennylane/workflow/qnode.py
index 83a7e8de95b..fa3baab25a2 100644
--- a/pennylane/workflow/qnode.py
+++ b/pennylane/workflow/qnode.py
@@ -25,8 +25,8 @@
 
 import pennylane as qml
 from pennylane import Device
-from pennylane.logging import debug_logger
 from pennylane.debugging import pldb_device_manager
+from pennylane.logging import debug_logger
 from pennylane.measurements import CountsMP, MidMeasureMP, Shots
 from pennylane.tape import QuantumScript, QuantumTape
 
diff --git a/tests/test_debugging.py b/tests/test_debugging.py
index e9970dc02b0..60cd4d737bb 100644
--- a/tests/test_debugging.py
+++ b/tests/test_debugging.py
@@ -444,7 +444,7 @@ def my_circ():
 
     def test_add_device(self):
         """Test that we can add a device to the global active device list."""
-        assert getattr(PLDB, "_PLDB__active_dev") is None
+        assert not PLDB.has_active_dev()
 
         dev1, dev2, dev3 = (
             qml.device("default.qubit", wires=3),
@@ -453,12 +453,12 @@ def test_add_device(self):
         )
 
         PLDB.add_device(dev1)
-        assert getattr(PLDB, "_PLDB__active_dev") == dev1
+        assert PLDB.get_active_device == dev1
 
         PLDB.add_device(dev2)  # overwrites dev1
         PLDB.add_device(dev3)  # overwrites dev2
 
-        assert getattr(PLDB, "_PLDB__active_dev") == dev3
+        assert PLDB.get_active_device == dev3
 
         PLDB.reset_active_dev()  # clean up the debugger active devices
 
@@ -471,17 +471,13 @@ def test_add_device(self):
     def test_get_active_device(self, device_name):
         """Test that we can access the active device."""
         dev = qml.device(device_name, wires=2)
-        PLDB.add_device(dev)
-
-        debugger_dev = PLDB.get_active_device()
-        assert debugger_dev is dev
-
-        PLDB.reset_active_dev()
+        with pldb_device_manager(dev) as _:
+            assert PLDB.get_active_device() is dev
 
     def test_get_active_device_error_when_no_active_device(self):
         """Test that an error is raised if we try to get
         the active device when there are no active devices."""
-        assert getattr(PLDB, "_PLDB__active_dev") == None
+        assert not PLDB.has_active_dev()
 
         with pytest.raises(RuntimeError, match="No active device to get"):
             _ = PLDB.get_active_device()
@@ -491,33 +487,33 @@ def test_reset_active_device(self, device_name):
         """Test that we can rest the global active device list."""
         dev = qml.device(device_name, wires=2)
         PLDB.add_device(dev)
-        assert getattr(PLDB, "_PLDB__active_dev") == dev
+        assert PLDB.get_active_device() == dev
 
         PLDB.reset_active_dev()
-        assert getattr(PLDB, "_PLDB__active_dev") == None
+        assert not PLDB.has_active_dev()
 
-    def test_is_active_device(self):
+    def test_has_active_device(self):
         """Test that we can determine if there is an active device."""
         assert getattr(PLDB, "_PLDB__active_dev") == None
 
         dev = qml.device("default.qubit")
         PLDB.add_device(dev)
-        assert PLDB.is_active_dev() is True
+        assert PLDB.has_active_dev()
 
         PLDB.reset_active_dev()
-        assert PLDB.is_active_dev() is False
+        assert not PLDB.has_active_dev()
 
 
 @pytest.mark.parametrize("device_name", ("default.qubit", "lightning.qubit"))
 def test_pldb_device_manager(device_name):
     """Test that the context manager works as expected."""
-    assert getattr(PLDB, "_PLDB__active_dev") == None
+    assert not PLDB.has_active_dev()
     dev = qml.device(device_name, wires=2)
 
     with pldb_device_manager(dev) as _:
-        assert getattr(PLDB, "_PLDB__active_dev") == dev
+        assert PLDB.get_active_device() == dev
 
-    assert getattr(PLDB, "_PLDB__active_dev") == None
+    assert not PLDB.has_active_dev()
 
 
 @patch.object(PLDB, "set_trace")

From 84890ffc6aad1fd9302e4550c112033ce2cde8a8 Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Wed, 5 Jun 2024 15:32:45 -0400
Subject: [PATCH 17/24] format tests

---
 tests/test_debugging.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_debugging.py b/tests/test_debugging.py
index 60cd4d737bb..c9cb2d41eb8 100644
--- a/tests/test_debugging.py
+++ b/tests/test_debugging.py
@@ -494,7 +494,7 @@ def test_reset_active_device(self, device_name):
 
     def test_has_active_device(self):
         """Test that we can determine if there is an active device."""
-        assert getattr(PLDB, "_PLDB__active_dev") == None
+        assert getattr(PLDB, "_PLDB__active_dev") is None
 
         dev = qml.device("default.qubit")
         PLDB.add_device(dev)

From 1d03db1b6a54de85edb3f4368d78889681a517ce Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Wed, 5 Jun 2024 16:36:35 -0400
Subject: [PATCH 18/24] fix bugs

---
 pennylane/debugging.py  | 21 +++++++++++++++------
 tests/test_debugging.py |  4 ++--
 2 files changed, 17 insertions(+), 8 deletions(-)

diff --git a/pennylane/debugging.py b/pennylane/debugging.py
index 136847a97e8..9db341e5eee 100644
--- a/pennylane/debugging.py
+++ b/pennylane/debugging.py
@@ -112,17 +112,26 @@ def get_snapshots(*args, **kwargs):
 
 
 class PLDB(pdb.Pdb):
-    """Custom debugging class integrated with Pdb."""
+    """Custom debugging class integrated with Pdb.
 
-    __active_dev: qml.devices.Device | None = None
+    This class is responsible for storing and updating a global device to be
+    used for executing quantum circuits while in debugging context. The core
+    debugger functionality is inherited from the native Python debugger (Pdb).
+
+    This class is not directly user-facing, but is interfaced with the
+    ``qml.breakpoint()`` function and ``pldb_device_manager`` context manager.
+    The former is responsible for launching the debugger prompt and the latter
+    is responsible with extracting and storing the ``qnode.device``.
+
+    The device information is used for validation checks and to execute measurements.
+    """
+
+    __active_dev = None
 
     def __init__(self, *args, **kwargs):
         """Initialize the debugger, and set custom prompt string."""
         super().__init__(*args, **kwargs)
-
-    @property
-    def prompt(self):
-        return "[pldb]: "
+        self.prompt = "[pldb]: "
 
     @classmethod
     def valid_context(cls):
diff --git a/tests/test_debugging.py b/tests/test_debugging.py
index c9cb2d41eb8..0fd21099202 100644
--- a/tests/test_debugging.py
+++ b/tests/test_debugging.py
@@ -453,12 +453,12 @@ def test_add_device(self):
         )
 
         PLDB.add_device(dev1)
-        assert PLDB.get_active_device == dev1
+        assert PLDB.get_active_device() == dev1
 
         PLDB.add_device(dev2)  # overwrites dev1
         PLDB.add_device(dev3)  # overwrites dev2
 
-        assert PLDB.get_active_device == dev3
+        assert PLDB.get_active_device() == dev3
 
         PLDB.reset_active_dev()  # clean up the debugger active devices
 

From 7ca926e568a103e6688f8f21754c311e0bf59872 Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Wed, 5 Jun 2024 16:57:02 -0400
Subject: [PATCH 19/24] missing test, [skip ci]

---
 tests/test_debugging.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/test_debugging.py b/tests/test_debugging.py
index e4b6f8024bc..d9f4460b59a 100644
--- a/tests/test_debugging.py
+++ b/tests/test_debugging.py
@@ -562,7 +562,7 @@ def test_tape():
     assert qml.equal(expected_tape, executed_tape)
 
 
-@pytest.mark.parametrize("measurement_process", (qml.expval(qml.Z(0)), qml.state()))
+@pytest.mark.parametrize("measurement_process", (qml.expval(qml.Z(0)), qml.state(), qml.probs()))
 @patch.object(PLDB, "_execute")
 def test_measure(mock_method, measurement_process):
     """Test that the private measure function doesn't modify the active queue"""

From 246a12b98cd23edf96a05e46670b996c17de90ef Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Thu, 6 Jun 2024 09:53:16 -0400
Subject: [PATCH 20/24] format, [skip ci]

---
 pennylane/debugging.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/pennylane/debugging.py b/pennylane/debugging.py
index b6446266d5b..d1203d9c290 100644
--- a/pennylane/debugging.py
+++ b/pennylane/debugging.py
@@ -14,10 +14,9 @@
 """
 This module contains functionality for debugging quantum programs on simulator devices.
 """
-
+import copy
 import pdb
 import sys
-import copy
 from contextlib import contextmanager
 
 import pennylane as qml

From 1803c0b09d3cd0c3407389c3966781a9eac3d737 Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Thu, 6 Jun 2024 13:59:25 -0400
Subject: [PATCH 21/24] address code review comments

---
 pennylane/debugging.py | 18 +++++++++---------
 1 file changed, 9 insertions(+), 9 deletions(-)

diff --git a/pennylane/debugging.py b/pennylane/debugging.py
index d1203d9c290..789661a7a9c 100644
--- a/pennylane/debugging.py
+++ b/pennylane/debugging.py
@@ -232,7 +232,7 @@ def state():
     """Compute the state of the quantum circuit.
 
     Returns:
-        Array(complex): Quantum state of the circuit.
+        Array(complex): quantum state of the circuit.
     """
     with qml.queuing.QueuingManager.stop_recording():
         m = qml.state()
@@ -244,10 +244,10 @@ def expval(op):
     """Compute the expectation value of an observable.
 
     Args:
-        op (Operator): The observable to compute the expectation value for.
+        op (Operator): the observable to compute the expectation value for
 
     Returns:
-        complex: Quantum state of the circuit.
+        complex: expectation value of the operator
     """
 
     qml.queuing.QueuingManager.active_context().remove(op)  # ensure we didn't accidentally queue op
@@ -262,12 +262,12 @@ def probs(wires=None, op=None):
     """Compute the probability distribution for the state.
     Args:
         wires (Union[Iterable, int, str, list]): the wires the operation acts on
-        op (Union[Observable, MeasurementValue]): Observable (with a ``diagonalizing_gates``
+        op (Union[Observable, MeasurementValue]): observable (with a ``diagonalizing_gates``
             attribute) that rotates the computational basis, or a  ``MeasurementValue``
             corresponding to mid-circuit measurements.
 
     Returns:
-        Array(float): The probability distribution of the bitstrings for the wires.
+        Array(float): the probability distribution of the bitstrings for the wires
     """
     if op:
         qml.queuing.QueuingManager.active_context().remove(
@@ -284,10 +284,10 @@ def _measure(measurement):
     """Perform the measurement.
 
     Args:
-        measurement (MeasurementProcess): The type of measurement to be performed
+        measurement (MeasurementProcess): the type of measurement to be performed
 
     Returns:
-        tuple(complex): Results from the measurement
+        tuple(complex): results from the measurement
     """
     active_queue = qml.queuing.QueuingManager.active_context()
     copied_queue = copy.deepcopy(active_queue)
@@ -298,10 +298,10 @@ def _measure(measurement):
 
 
 def tape():
-    """Access the tape of the quantum circuit.
+    """Access the quantum tape of the circuit.
 
     Returns:
-        QuantumScript: The quantum script representing the circuit.
+        QuantumScript: the quantum tape representing the circuit
     """
     active_queue = qml.queuing.QueuingManager.active_context()
     return qml.tape.QuantumScript.from_queue(active_queue)

From d554e2344c8014ea21256d3cd77b7223bc69cb0e Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Fri, 7 Jun 2024 09:55:21 -0400
Subject: [PATCH 22/24] Apply suggestions from code review

Co-authored-by: Utkarsh <utkarshazad98@gmail.com>
---
 pennylane/debugging.py | 10 ++++------
 1 file changed, 4 insertions(+), 6 deletions(-)

diff --git a/pennylane/debugging.py b/pennylane/debugging.py
index 789661a7a9c..f2456ad2f32 100644
--- a/pennylane/debugging.py
+++ b/pennylane/debugging.py
@@ -192,11 +192,9 @@ def _execute(cls, batch_tapes):
         """Execute tape on the active device"""
         dev = cls.get_active_device()
 
-        valid_batch, _ = (
-            (batch_tapes, None)
-            if not dev.wires
-            else qml.devices.preprocess.validate_device_wires(batch_tapes, wires=dev.wires)
-        )
+        valid_batch = batch_tapes
+        if dev.wires:
+            valid_batch = qml.devices.preprocess.validate_device_wires(batch_tapes, wires=dev.wires)[0]
 
         program, new_config = dev.preprocess()
         new_batch, fn = program(valid_batch)
@@ -263,7 +261,7 @@ def probs(wires=None, op=None):
     Args:
         wires (Union[Iterable, int, str, list]): the wires the operation acts on
         op (Union[Observable, MeasurementValue]): observable (with a ``diagonalizing_gates``
-            attribute) that rotates the computational basis, or a  ``MeasurementValue``
+            attribute) that rotates the computational basis, or a ``MeasurementValue``
             corresponding to mid-circuit measurements.
 
     Returns:

From ebfbc7f0057eb97bc7af455a89a596e6b91fbc44 Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Fri, 7 Jun 2024 10:05:31 -0400
Subject: [PATCH 23/24] address code review comments

---
 pennylane/debugging.py  |  4 +++-
 tests/test_debugging.py | 10 +++++-----
 2 files changed, 8 insertions(+), 6 deletions(-)

diff --git a/pennylane/debugging.py b/pennylane/debugging.py
index f2456ad2f32..097aea6e274 100644
--- a/pennylane/debugging.py
+++ b/pennylane/debugging.py
@@ -194,7 +194,9 @@ def _execute(cls, batch_tapes):
 
         valid_batch = batch_tapes
         if dev.wires:
-            valid_batch = qml.devices.preprocess.validate_device_wires(batch_tapes, wires=dev.wires)[0]
+            valid_batch = qml.devices.preprocess.validate_device_wires(
+                batch_tapes, wires=dev.wires
+            )[0]
 
         program, new_config = dev.preprocess()
         new_batch, fn = program(valid_batch)
diff --git a/tests/test_debugging.py b/tests/test_debugging.py
index d9f4460b59a..40208e1e808 100644
--- a/tests/test_debugging.py
+++ b/tests/test_debugging.py
@@ -569,7 +569,7 @@ def test_measure(mock_method, measurement_process):
     with qml.queuing.AnnotatedQueue() as queue:
         ops = [qml.X(0), qml.Y(1), qml.Z(0)] + [qml.Hadamard(i) for i in range(3)]
         measurements = [qml.expval(qml.X(2)), qml.state(), qml.probs(), qml.var(qml.Z(3))]
-        _ = qml.debugging._measure(measurement_process)
+        qml.debugging._measure(measurement_process)
 
     executed_tape = qml.tape.QuantumScript.from_queue(queue)
     expected_tape = qml.tape.QuantumScript(ops, measurements)
@@ -592,7 +592,7 @@ def test_state(_mock_method):
         qml.RY(0.45, 2)
         qml.sample()
 
-        _ = qml.debugging.state()
+        qml.debugging.state()
 
     assert qml.state() not in queue
 
@@ -606,7 +606,7 @@ def test_expval(_mock_method):
             qml.RY(0.45, 2)
             qml.sample()
 
-            _ = qml.debugging.expval(op)
+            qml.debugging.expval(op)
 
         assert op not in queue
         assert qml.expval(op) not in queue
@@ -622,7 +622,7 @@ def test_probs_with_op(_mock_method):
             qml.RY(0.45, 2)
             qml.sample()
 
-            _ = qml.debugging.probs(op=op)
+            qml.debugging.probs(op=op)
 
         assert op not in queue
         assert qml.probs(op=op) not in queue
@@ -638,7 +638,7 @@ def test_probs_with_wires(_mock_method):
             qml.RY(0.45, 2)
             qml.sample()
 
-            _ = qml.debugging.probs(wires=wires)
+            qml.debugging.probs(wires=wires)
 
         assert qml.probs(wires=wires) not in queue
 

From aa7d354a7927cf0b29fdf7c7897e2b7f073c22ec Mon Sep 17 00:00:00 2001
From: Jay Soni <jbsoni@uwaterloo.ca>
Date: Fri, 7 Jun 2024 10:18:02 -0400
Subject: [PATCH 24/24] lint + format

---
 pennylane/debugging.py  | 4 ++--
 tests/test_debugging.py | 5 ++++-
 2 files changed, 6 insertions(+), 3 deletions(-)

diff --git a/pennylane/debugging.py b/pennylane/debugging.py
index 097aea6e274..21ff5be40f3 100644
--- a/pennylane/debugging.py
+++ b/pennylane/debugging.py
@@ -293,8 +293,8 @@ def _measure(measurement):
     copied_queue = copy.deepcopy(active_queue)
 
     copied_queue.append(measurement)
-    tape = qml.tape.QuantumScript.from_queue(copied_queue)
-    return PLDB._execute((tape,))
+    qtape = qml.tape.QuantumScript.from_queue(copied_queue)
+    return PLDB._execute((qtape,))  # pylint: disable=protected-access
 
 
 def tape():
diff --git a/tests/test_debugging.py b/tests/test_debugging.py
index 40208e1e808..30d49f6d387 100644
--- a/tests/test_debugging.py
+++ b/tests/test_debugging.py
@@ -551,7 +551,10 @@ def test_tape():
     """Test that we can access the tape from the active queue."""
     with qml.queuing.AnnotatedQueue() as queue:
         qml.X(0)
-        [qml.Hadamard(i) for i in range(3)]
+
+        for i in range(3):
+            qml.Hadamard(i)
+
         qml.Y(1)
         qml.Z(0)
         qml.expval(qml.Z(0))