Repository: PennyLaneAI/pennylane
Pull request #5768 (2024-05-30 14:57:08)

Quantum debugging (feature branch)
**Context:**
Feature branch for the new quantum debugging feature

github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5768?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.66%. Comparing base [(`f5b805b`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/f5b805b1c57b3deaaa22ba4b2166c4dd77bc500b?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`2c25488`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/2c25488771d92b278a60e142e3d260effbb63adc?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 6 commits behind head on master.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #5768      +/-   ##
==========================================
- Coverage   99.67%   99.66%   -0.02%     
==========================================
  Files         416      416              
  Lines       38686    38741      +55     
==========================================
+ Hits        38562    38612      +50     
- Misses        124      129       +5     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5768?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).


Patch:
diff --git a/pennylane/__init__.py b/pennylane/__init__.py
index 139ebc292..9e92671d2 100644
--- a/pennylane/__init__.py
+++ b/pennylane/__init__.py
@@ -123,7 +123,7 @@ from pennylane.ops.functions import (
 )
 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 abfb4c5d2..0d91bb827 100644
--- a/pennylane/debugging.py
+++ b/pennylane/debugging.py
@@ -14,6 +14,11 @@
 """
 This module contains functionality for debugging quantum programs on simulator devices.
 """
+import copy
+import pdb
+import sys
+from contextlib import contextmanager
+
 import pennylane as qml
 from pennylane import DeviceError
 
@@ -105,3 +110,198 @@ def snapshots(qnode):
         return dbg.snapshots
 
     return get_snapshots
+
+
+class PLDB(pdb.Pdb):
+    """Custom debugging class integrated with Pdb.
+
+    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)
+        self.prompt = "[pldb]: "
+
+    @classmethod
+    def valid_context(cls):
+        """Determine if the debugger is called in a valid context.
+
+        Raises:
+            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.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"):
+            raise TypeError("Breakpoints not supported on this device")
+
+    @classmethod
+    def add_device(cls, dev):
+        """Update the global active device.
+
+        Args:
+            dev (Union[Device, "qml.devices.Device"]): the active device
+        """
+        cls.__active_dev = dev
+
+    @classmethod
+    def get_active_device(cls):
+        """Return the active device.
+
+        Raises:
+            RuntimeError: No active device to get
+
+        Returns:
+            Union[Device, "qml.devices.Device"]: The active device
+        """
+        if not cls.has_active_dev():
+            raise RuntimeError("No active device to get")
+
+        return cls.__active_dev
+
+    @classmethod
+    def has_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):
+        """Reset the global active device variable to None."""
+        cls.__active_dev = None
+
+    @classmethod
+    def _execute(cls, batch_tapes):
+        """Execute tape on the active device"""
+        dev = cls.get_active_device()
+
+        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)
+
+        # TODO: remove [0] index once compatible with transforms
+        return fn(dev.execute(new_batch, new_config))[0]
+
+
+@contextmanager
+def pldb_device_manager(device):
+    """Context manager to automatically set and reset active
+    device on the Pennylane Debugger (PLDB).
+
+    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(skip=["pennylane.*"])  # skip internals when stepping through trace
+    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: expectation value of the operator
+    """
+
+    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 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.
+
+    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)
+    qtape = qml.tape.QuantumScript.from_queue(copied_queue)
+    return PLDB._execute((qtape,))  # pylint: disable=protected-access
+
+
+def tape():
+    """Access the quantum tape of the circuit.
+
+    Returns:
+        QuantumScript: the quantum tape representing the circuit
+    """
+    active_queue = qml.queuing.QueuingManager.active_context()
+    return qml.tape.QuantumScript.from_queue(active_queue)
diff --git a/pennylane/workflow/qnode.py b/pennylane/workflow/qnode.py
index e0fd75458..fa3baab25 100644
--- a/pennylane/workflow/qnode.py
+++ b/pennylane/workflow/qnode.py
@@ -25,6 +25,7 @@ from typing import Optional, Union
 
 import pennylane as qml
 from pennylane import Device
+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
@@ -958,8 +959,12 @@ class QNode:
         if old_interface == "auto":
             self.interface = qml.math.get_interface(*args, *list(kwargs.values()))
 
-        with qml.queuing.AnnotatedQueue() as q:
-            self._qfunc_output = self.func(*args, **kwargs)
+        # Before constructing the tape, we pass the device to the
+        # debugger to ensure they are compatible if there are any
+        # breakpoints in the circuit
+        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)
 


Test patch:
diff --git a/tests/test_debugging.py b/tests/test_debugging.py
index 0434fe22a..30d49f6d3 100644
--- a/tests/test_debugging.py
+++ b/tests/test_debugging.py
@@ -14,10 +14,14 @@
 """
 Unit tests for the debugging module.
 """
+from unittest.mock import patch
+
 import numpy as np
 import pytest
 
 import pennylane as qml
+from pennylane import numpy as qnp
+from pennylane.debugging import PLDB, pldb_device_manager
 
 
 class TestSnapshot:
@@ -387,3 +391,303 @@ class TestSnapshot:
             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") is None
+
+    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(
+            RuntimeError, match="Can't call breakpoint outside of a qnode execution"
+        ):
+            with qml.queuing.AnnotatedQueue() as _:
+                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(
+            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 with an incompatible device."""
+        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="Breakpoints not supported on this device"):
+            _ = 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 not PLDB.has_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 PLDB.get_active_device() == dev1
+
+        PLDB.add_device(dev2)  # overwrites dev1
+        PLDB.add_device(dev3)  # overwrites dev2
+
+        assert PLDB.get_active_device() == dev3
+
+        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 access the active device."""
+        dev = qml.device(device_name, wires=2)
+        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 not PLDB.has_active_dev()
+
+        with pytest.raises(RuntimeError, 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 PLDB.get_active_device() == dev
+
+        PLDB.reset_active_dev()
+        assert not PLDB.has_active_dev()
+
+    def test_has_active_device(self):
+        """Test that we can determine if there is an active device."""
+        assert getattr(PLDB, "_PLDB__active_dev") is None
+
+        dev = qml.device("default.qubit")
+        PLDB.add_device(dev)
+        assert PLDB.has_active_dev()
+
+        PLDB.reset_active_dev()
+        assert not PLDB.has_active_dev()
+
+    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))],
+        ),
+        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])],
+        ),
+        qml.tape.QuantumScript(
+            ops=[qml.Hadamard(0)],
+            measurements=[qml.state()],
+        ),  # Test that state expands to number of device wires
+    )
+
+    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]),
+        qnp.array([1 / qnp.sqrt(2), 0, 1 / qnp.sqrt(2), 0], dtype=complex),
+    )
+
+    @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)
+        PLDB.reset_active_dev()
+
+
+def test_tape():
+    """Test that we can access the tape from the active queue."""
+    with qml.queuing.AnnotatedQueue() as queue:
+        qml.X(0)
+
+        for i in range(3):
+            qml.Hadamard(i)
+
+        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(), 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"""
+    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.sample()
+
+        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.sample()
+
+            qml.debugging.expval(op)
+
+        assert op not in queue
+        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."""
+    assert not PLDB.has_active_dev()
+    dev = qml.device(device_name, wires=2)
+
+    with pldb_device_manager(dev) as _:
+        assert PLDB.get_active_device() == dev
+
+    assert not PLDB.has_active_dev()
+
+
+@patch.object(PLDB, "set_trace")
+def test_breakpoint_integration(mock_method):
+    """Test that qml.breakpoint behaves as expected"""
+    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="Breakpoints not supported on this device"):
+        _ = my_circ()
+
+    mock_method.assert_not_called()  # Error was raised before we triggered breakpoint