Add support for qjit'd code to qml.specs()

doc/releases/changelog-dev.md

@@ -3,13 +3,37 @@
 
 <h3>New features since last release</h3>
 
+* The `qml.specs` function can now support code that has been compiled with `qml.qjit`.
+  This new feature only supports getting circuit specs at the `device` level.
+  [(#8202)](https://github.com/PennyLaneAI/pennylane/pull/8202)
+
+  ```python
+  @qml.qjit
+  @qml.qnode(qml.device("lightning.qubit", wires=2))
+  def circuit():
+      qml.Hadamard(wires=0)
+      qml.CNOT(wires=[0, 1])
+      return qml.expval(qml.Z(0) @ qml.Z(1))
+
+  print(qml.specs(circuit, level="device")())
+  ```
+  ```
+  {'resources': Resources(num_wires=2, num_gates=2, gate_types=defaultdict(<class 'int'>, {'CNOT': 1, 'Hadamard': 1}), gate_sizes=defaultdict(<class 'int'>, {2: 1, 1: 1}), depth=2, shots=Shots(total_shots=None, shot_vector=())),
+   'num_device_wires': 2,
+   'device_name': 'lightning.qubit',
+   'level': 'device',
+   'gradient_options': {},
+   'interface': 'auto',
+   'diff_method': 'best'}
+  ```
+
 * The `qml.sample` function can now receive an optional `dtype` parameter 
   which sets the type and precision of the samples returned by this measurement process.
   [(#8189)](https://github.com/PennyLaneAI/pennylane/pull/8189)
 
 * The Resource estimation toolkit was upgraded and has migrated from
   :mod:`~.labs` to PennyLane as the :mod:`~.estimator` module.
-  
+
   * The `qml.estimator.WireResourceManager`, `qml.estimator.Allocate`, and `qml.estimator.Deallocate`
     classes were added to track auxiliary wires for resource estimation.
     [(#8203)](https://github.com/PennyLaneAI/pennylane/pull/8203)

pennylane/devices/null_qubit.py

@@ -18,10 +18,7 @@
 
 
 import inspect
-import json
 import logging
-import time
-from collections import defaultdict
 from dataclasses import replace
 from functools import lru_cache, singledispatch
 from numbers import Number
@@ -40,7 +37,6 @@
     Shots,
     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
@@ -52,8 +48,6 @@
 logger = logging.getLogger(__name__)
 logger.addHandler(logging.NullHandler())
 
-RESOURCES_FNAME_PREFIX = "__pennylane_resources_data_"
-
 
 @singledispatch
 def zero_measurement(
@@ -135,49 +129,6 @@ def _interface(config: ExecutionConfig):
     return config.interface.get_like() if config.gradient_method == "backprop" else "numpy"
 
 
-def _simulate_resource_use(circuit, outfile):
-    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 behaviour of catalyst resource tracking.
-    #  In the future it may be better to return this information in a more structured way.
-    json.dump(
-        {
-            "num_wires": num_wires,
-            "num_gates": sum(gate_types.values()),
-            "gate_types": gate_types,
-        },
-        outfile,
-    )
-
-
 @simulator_tracking
 @single_tape_support
 class NullQubit(Device):
@@ -190,10 +141,11 @@ 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 and save them to a JSON file.
-            The filename will match ``__pennylane_resources_data_*.json`` where the wildcard (asterisk) is replaced by
-            the timestamp of when execution began in nanoseconds since Unix EPOCH. This argument is experimental and
-            subject to change.
+        track_resources (bool): If True, turn on Catalyst device resource tracking.
+        resources_filename (string): If set, the static filename to use when saving resource data.
+            If not set, the filename will match ``__pennylane_resources_data_*`` where the wildcard (asterisk)
+            is replaced by the timestamp of when execution began in nanoseconds since Unix EPOCH.
+        compute_depth (bool): If True, compute the circuit depth as part of resource tracking.
 
     **Example:**
 
@@ -278,24 +230,34 @@ def name(self):
         """The name of the device."""
         return "null.qubit"
 
-    def __init__(self, wires=None, shots=None, track_resources=False) -> None:
+    # pylint: disable=too-many-arguments
+    def __init__(
+        self,
+        wires=None,
+        shots=None,
+        track_resources=False,
+        resources_filename=None,
+        compute_depth=None,
+        target_device=None,
+    ) -> None:
         super().__init__(wires=wires, shots=shots)
         self._debugger = None
-        self._track_resources = track_resources
+        self._target_device = target_device
 
         # this is required by Catalyst to toggle the tracker at runtime
         self.device_kwargs = {"track_resources": track_resources}
+        if resources_filename is not None:
+            self.device_kwargs["resources_filename"] = resources_filename
+        if compute_depth is not None:
+            self.device_kwargs["compute_depth"] = compute_depth
+
+        if target_device is not None:
+            self.config_filepath = target_device.config_filepath
 
     def _simulate(self, circuit, interface):
         num_device_wires = len(self.wires) if self.wires else len(circuit.wires)
         results = []
 
-        if self._track_resources:
-            timestamp = int(time.time() * 1e9)  # nanoseconds since epoch
-            resources_fname = f"{RESOURCES_FNAME_PREFIX}{timestamp}.json"
-            with open(resources_fname, "x", encoding="utf-8") as f:
-                _simulate_resource_use(circuit, f)
-
         for s in circuit.shots or [None]:
             r = tuple(
                 zero_measurement(mp, num_device_wires, s, circuit.batch_size, interface)
@@ -347,8 +309,13 @@ def preprocess(
         if execution_config is None:
             execution_config = ExecutionConfig()
 
-        target = DefaultQubit(wires=self.wires)
-        program = target.preprocess_transforms(execution_config)
+        if self._target_device is None:
+            target = DefaultQubit(wires=self.wires)
+        else:
+            target = self._target_device
+
+        program, _ = target.preprocess(execution_config)
+
         for t in program:
             if t.transform == decompose.transform:
                 original_stopping_condition = t.kwargs["stopping_condition"]