Repository: PennyLaneAI/pennylane
Pull request #7226 (2025-04-07 19:25:32)

Add resource tracking to null.qubit
**Context:**
For benchmarking, we need a method of precise resource tracking. This method has to be compatible with code `@qjit` compiled using Catalyst and should be able to track precise resources after various transformation passes have already occurred.

See also: https://github.com/PennyLaneAI/pennylane-benchmarks/pull/76

**Description of the Change:**
Adds a new attribute to `null.qubit`, as well as an optional argument to its `__init__` function. When this option is toggled on, `null.qubit` will track precise gate counts when "executing" a circuit.

**Benefits:**
It is now possible to track precise gate counts (with or without gate decomposition or other transformation passes) for measuring circuit complexity.

**Possible Drawbacks:**
To be mutually compatible with the catalyst implementation, this currently just saves the resources JSON to a file as opposed to neatly returning an object. In the future, a better solution may be desired.
This implementation causes `null.qubit` to have a different capability that other devices do not have. If this feature is useful elsewhere, it might be wise to move resource tracking to the base Device class.

**Related GitHub Issues:**
[sc-88213]

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.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 0d3749676..8b5b6760c 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -29,7 +29,6 @@
   [0.87758256+0.j 0.47942554+0.j 0.        +0.j 0.        +0.j
    0.        +0.j 0.        +0.j 0.        +0.j 0.        +0.j]
   ```
-  
 
 * The transform `convert_to_mbqc_gateset` is added to the `ftqc` module to convert arbitrary 
   circuits to a limited gate-set that can be translated to the MBQC formalism.
@@ -220,6 +219,9 @@
 
 <h3>Internal changes ⚙️</h3>
 
+* `null.qubit` can now support an optional `track_resources` argument which allows it to record which gates are executed.
+  [(#7226)](https://github.com/PennyLaneAI/pennylane/pull/7226)
+
 * A new internal module, `qml.concurrency`, is added to support internal use of multiprocess and multithreaded execution of workloads. This also migrates the use of `concurrent.futures` in `default.qubit` to this new design.
   [(#7303)](https://github.com/PennyLaneAI/pennylane/pull/7303)
 
@@ -313,5 +315,6 @@ Lillian Frederiksen,
 Pietropaolo Frisoni,
 Korbinian Kottmann,
 Christina Lee,
+Andrija Paurevic,
 Lee J. O'Riordan,
-Andrija Paurevic
+Jake Zaia
diff --git a/pennylane/devices/null_qubit.py b/pennylane/devices/null_qubit.py
index 5ced33798..4e172baa0 100644
--- a/pennylane/devices/null_qubit.py
+++ b/pennylane/devices/null_qubit.py
@@ -18,7 +18,9 @@ benchmarking PennyLane's auxiliary functionality outside direct circuit evaluati
 # pylint:disable=unused-argument
 
 import inspect
+import json
 import logging
+from collections import defaultdict
 from dataclasses import replace
 from functools import lru_cache, singledispatch
 from numbers import Number
@@ -37,6 +39,7 @@ from pennylane.measurements import (
     ProbabilityMP,
     StateMP,
 )
+from pennylane.ops.op_math import Adjoint, Controlled, ControlledOp
 from pennylane.tape import QuantumScriptOrBatch
 from pennylane.transforms.core import TransformProgram
 from pennylane.typing import Result, ResultBatch
@@ -129,6 +132,50 @@ def _interface(config: ExecutionConfig):
     return config.interface.get_like() if config.gradient_method == "backprop" else "numpy"
 
 
+def _simulate_resource_use(circuit, resources_fname="__pennylane_resources_data.json"):
+    num_wires = len(circuit.wires)
+    gate_types = defaultdict(int)
+
+    for op in circuit.operations:
+        controls = 0
+        adj = False
+
+        while hasattr(op, "base"):
+            if type(op) in (Controlled, ControlledOp):
+                # Don't check this with `isinstance` to avoid unrolling ops like CNOT
+                controls += len(op.control_wires)
+            elif isinstance(op, Adjoint):
+                adj = not adj
+            else:
+                break  # Certain gates have "base" but shouldn't be broken down (like CNOT)
+            # NOTE: Pow, Exp, Add, etc. intentionally not handled to be compatible with Catalyst
+            op = op.base
+
+        # Barrier is not a gate and takes no resources, so we skip it
+        if op.name == "Barrier":
+            # (this confuses codecov, despite being tested)
+            continue  # pragma: no cover
+
+        name = op.name
+        if adj:
+            name = f"Adj({name})"
+        if controls:
+            name = f"{controls if controls > 1 else ''}C({name})"
+
+        gate_types[name] += 1
+    # NOTE: For now, this information is being printed to match the behavior of catalyst resource tracking.
+    #  In the future it may be better to return this information in a more structured way.
+    with open(resources_fname, "w") as f:
+        json.dump(
+            {
+                "num_wires": num_wires,
+                "num_gates": sum(gate_types.values()),
+                "gate_types": gate_types,
+            },
+            f,
+        )
+
+
 @simulator_tracking
 @single_tape_support
 class NullQubit(Device):
@@ -141,7 +188,7 @@ class NullQubit(Device):
             (``['aux_wire', 'q1', 'q2']``). Default ``None`` if not specified.
         shots (int, Sequence[int], Sequence[Union[int, Sequence[int]]]): The default number of shots
             to use in executions involving this device.
-
+        track_resources (bool): If ``True``, track the number of resources used by the device. This argument is experimental and subject to change.
     **Example:**
 
     .. code-block:: python
@@ -225,13 +272,18 @@ class NullQubit(Device):
         """The name of the device."""
         return "null.qubit"
 
-    def __init__(self, wires=None, shots=None) -> None:
+    def __init__(self, wires=None, shots=None, track_resources=False) -> None:
         super().__init__(wires=wires, shots=shots)
         self._debugger = None
+        self._track_resources = track_resources
 
     def _simulate(self, circuit, interface):
         num_device_wires = len(self.wires) if self.wires else len(circuit.wires)
         results = []
+
+        if self._track_resources:
+            _simulate_resource_use(circuit)
+
         for s in circuit.shots or [None]:
             r = tuple(
                 zero_measurement(mp, num_device_wires, s, circuit.batch_size, interface)


Test patch:
diff --git a/tests/devices/test_null_qubit.py b/tests/devices/test_null_qubit.py
index 577cd4efe..e0570e569 100644
--- a/tests/devices/test_null_qubit.py
+++ b/tests/devices/test_null_qubit.py
@@ -13,6 +13,8 @@
 # limitations under the License.
 """Tests for null.qubit."""
 
+import json
+import os
 from collections import defaultdict as dd
 
 import numpy as np
@@ -61,6 +63,73 @@ def test_debugger_attribute():
     assert dev._debugger is None
 
 
+def test_resource_tracking_attribute():
+    """Test NullQubit track_resources attribute"""
+    # pylint: disable=protected-access
+    assert NullQubit()._track_resources is False
+    assert NullQubit(track_resources=True)._track_resources is True
+
+    dev = NullQubit(track_resources=True)
+
+    def small_circ(params):
+        qml.X(0)
+        qml.H(0)
+
+        qml.Barrier()
+
+        # Add a more complex operation to check that the innermost operation is counted
+        op = qml.T(0)
+        op = qml.adjoint(op)
+        op = qml.ctrl(op, control=1, control_values=[1])
+
+        qml.ctrl(qml.S(0), control=[1, 2], control_values=[1, 1])
+
+        qml.CNOT([0, 1])
+        qml.Barrier()
+
+        qml.ctrl(qml.IsingXX(0, [0, 1]), control=2, control_values=[1])
+        qml.adjoint(qml.S(0))
+
+        qml.RX(params[0], wires=0)
+        qml.RX(params[0] * 2, wires=1)
+
+        return qml.expval(qml.PauliZ(0))
+
+    qnode = qml.QNode(small_circ, dev, diff_method="backprop")
+
+    inputs = qml.numpy.array([0.5])
+
+    qnode(inputs)
+
+    # Check that resource tracking doesn't interfere with backprop
+    assert qml.grad(qnode)(inputs) == 0
+
+    RESOURCES_FNAME = "__pennylane_resources_data.json"
+    assert os.path.exists(RESOURCES_FNAME)
+
+    with open(RESOURCES_FNAME, "r") as f:
+        stats = f.read()
+
+    os.remove(RESOURCES_FNAME)
+
+    assert stats == json.dumps(
+        {
+            "num_wires": 3,
+            "num_gates": 9,
+            "gate_types": {
+                "PauliX": 1,
+                "Hadamard": 1,
+                "C(Adj(T))": 1,
+                "2C(S)": 1,
+                "CNOT": 1,
+                "C(IsingXX)": 1,
+                "Adj(S)": 1,
+                "RX": 2,
+            },
+        }
+    )
+
+
 @pytest.mark.parametrize("shots", (None, 10))
 def test_supports_operator_without_decomp(shots):
     """Test that null.qubit automatically supports any operation without a decomposition."""