Lazy loading of Quantumscript properties (#5696)

**Context:**
Eager initialization of parameters inside the constructor of
`QuantumScript` induces a high classical overhead.

**Description of the Change:**
The following attributes of `QuantumScript` are no longer initialized in
the constructor, but only when needed:

- par_info
- wires, num_wires
- obs_sharing_wires, obs_sharing_wires_id

`par_info`, `obs_sharing_wires`, and `obs_sharing_wires_id` are now
public attributes as they were used in several files outside the
`QuantumScript` class. Examples of this behaviour are:

- pennylane/_qubit_device.py
- pennylane/pennylane/circuit_graph.py
- pennylane/pennylane/fourier/qnode_spectrum.py
- pennylane/pennylane/gradients/adjoint_metric_tensor.py
- pennylane/pennylane/gradients/gradient_transform.py
- pennylane/pennylane/transforms/tape_expand.py


The most expensive attributes: `par_info` and `wires ` are now a
`@cached_property` to avoid calculating them every time they are needed.

When used inside the `QuantumTape` context, the `_update()` method is
called to invalidate the cached value and allow for recalculation.


**Benefits:**
Reduce in the classical overhead of creating a `QuantumScript`. For
instance, for the following code:

```python
n_wires = 5
n_layers = 50

dev = qml.device('null.qubit')

@qml.qnode(dev, diff_method="parameter-shift")
def circuit(params):
    qml.StronglyEntanglingLayers(params, wires=range(n_wires))
    return qml.expval(qml.PauliZ(n_wires-1))

shape = qml.StronglyEntanglingLayers.shape(n_wires=n_wires, n_layers=n_layers)
rng = np.random.default_rng(seed=42)
params = np.array(rng.random(shape))
qml.grad(circuit)(params)
```
The `QuantumScript` creation time decreases from 62.51% of the whole
time, to 41.92%.

Related Shortcut Stories:
[sc-45973]

doc/releases/changelog-dev.md

@@ -102,6 +102,10 @@
 * Empty initialization of `PauliVSpace` is permitted.
   [(#5675)](https://github.com/PennyLaneAI/pennylane/pull/5675)
 
+* `QuantumScript` properties are only calculated when needed, instead of on initialization. This decreases the classical overhead by >20%.
+  `par_info`, `obs_sharing_wires`, and `obs_sharing_wires_id` are now public attributes.
+  [(#5696)](https://github.com/PennyLaneAI/pennylane/pull/5696)
+
 <h4>Community contributions 🥳</h4>
 
 * Implemented kwargs (`check_interface`, `check_trainability`, `rtol` and `atol`) support in `qml.equal` for the operators `Pow`, `Adjoint`, `Exp`, and `SProd`.
@@ -197,6 +201,7 @@ Astral Cai,
 Ahmed Darwish,
 Isaac De Vlugt,
 Pietropaolo Frisoni,
+Emiliano Godinez,
 Soran Jahangiri,
 Korbinian Kottmann,
 Christina Lee,

pennylane/_device.py

@@ -16,7 +16,6 @@
 """
 # pylint: disable=too-many-format-args, use-maxsplit-arg, protected-access
 import abc
-import copy
 import types
 import warnings
 from collections import OrderedDict
@@ -81,7 +80,7 @@ def _local_tape_expand(tape, depth, stop_at):
             if isinstance(obj, Operator):
                 if obj.has_decomposition:
                     with QueuingManager.stop_recording():
-                        obj = QuantumScript(obj.decomposition(), _update=False)
+                        obj = QuantumScript(obj.decomposition())
                 else:
                     new_queue.append(obj)
                     continue
@@ -94,11 +93,9 @@ def _local_tape_expand(tape, depth, stop_at):
 
     # preserves inheritance structure
     # if tape is a QuantumTape, returned object will be a quantum tape
-    new_tape = tape.__class__(new_ops, new_measurements, shots=tape.shots, _update=False)
+    new_tape = tape.__class__(new_ops, new_measurements, shots=tape.shots)
 
     # Update circuit info
-    new_tape.wires = copy.copy(tape.wires)
-    new_tape.num_wires = tape.num_wires
     new_tape._batch_size = tape._batch_size
     new_tape._output_dim = tape._output_dim
     return new_tape
@@ -675,9 +672,9 @@ def default_expand_fn(self, circuit, max_expansion=10):
         comp_basis_sampled_multi_measure = (
             len(circuit.measurements) > 1 and circuit.samples_computational_basis
         )
-        obs_on_same_wire = len(circuit._obs_sharing_wires) > 0 or comp_basis_sampled_multi_measure
+        obs_on_same_wire = len(circuit.obs_sharing_wires) > 0 or comp_basis_sampled_multi_measure
         obs_on_same_wire &= not any(
-            isinstance(o, (Hamiltonian, LinearCombination)) for o in circuit._obs_sharing_wires
+            isinstance(o, (Hamiltonian, LinearCombination)) for o in circuit.obs_sharing_wires
         )
         ops_not_supported = not all(self.stopping_condition(op) for op in circuit.operations)
 
@@ -688,7 +685,6 @@ def default_expand_fn(self, circuit, max_expansion=10):
             circuit = _local_tape_expand(
                 circuit, depth=max_expansion, stop_at=self.stopping_condition
             )
-            circuit._update()
 
         return circuit
 
@@ -780,7 +776,7 @@ def batch_transform(self, circuit: QuantumTape):
             circuits, hamiltonian_fn = qml.transforms.sum_expand(circuit)
 
         elif (
-            len(circuit._obs_sharing_wires) > 0
+            len(circuit.obs_sharing_wires) > 0
             and not hamiltonian_in_obs
             and all(
                 not isinstance(m, (SampleMP, ProbabilityMP, CountsMP)) for m in circuit.measurements

pennylane/_qubit_device.py

@@ -1715,7 +1715,7 @@ def adjoint_jacobian(
         trainable_params = []
         for k in tape.trainable_params:
             # pylint: disable=protected-access
-            mp_or_op = tape[tape._par_info[k]["op_idx"]]
+            mp_or_op = tape[tape.par_info[k]["op_idx"]]
             if isinstance(mp_or_op, MeasurementProcess):
                 warnings.warn(
                     "Differentiating with respect to the input parameters of "

pennylane/fourier/qnode_spectrum.py

@@ -409,7 +409,7 @@ def wrapper(*args, **kwargs):
         cjacs = jac_fn(*args, **kwargs)
         spectra = {}
         tape = qml.transforms.expand_multipar(qnode.qtape)
-        par_info = tape._par_info
+        par_info = tape.par_info
 
         # Iterate over jacobians per argument
         for jac_idx, cjac in enumerate(cjacs):

pennylane/gradients/adjoint_metric_tensor.py

@@ -40,7 +40,7 @@ def _group_operations(tape):
     ops = tape.operations
     # Find the indices of trainable operations in the tape operations list
     # pylint: disable=protected-access
-    trainable_par_info = [tape._par_info[i] for i in tape.trainable_params]
+    trainable_par_info = [tape.par_info[i] for i in tape.trainable_params]
     trainables = [info["op_idx"] for info in trainable_par_info]
     # Add the indices incremented by one to the trainable indices
     split_ids = list(chain.from_iterable([idx, idx + 1] for idx in trainables))

pennylane/gradients/gradient_transform.py

@@ -169,7 +169,7 @@ def _try_zero_grad_from_graph_or_get_grad_method(tape, param_index, use_graph=Tr
     """
 
     # pylint:disable=protected-access
-    par_info = tape._par_info[param_index]
+    par_info = tape.par_info[param_index]
 
     if use_graph:
         op_or_mp = tape[par_info["op_idx"]]

pennylane/gradients/hadamard_gradient.py

@@ -342,12 +342,14 @@ def _expval_hadamard_grad(tape, argnum, aux_wire):
                     measurements.append(qml.probs(op=obs_new))
 
             new_tape = qml.tape.QuantumScript(ops=ops, measurements=measurements, shots=tape.shots)
-
             _rotations, _measurements = qml.tape.tape.rotations_and_diagonal_measurements(new_tape)
-            # pylint: disable=protected-access
-            new_tape._ops = new_tape.operations + _rotations
-            new_tape._measurements = _measurements
-            new_tape._update()
+            new_ops = new_tape.operations + _rotations
+            new_tape = qml.tape.QuantumScript(
+                new_ops,
+                _measurements,
+                shots=new_tape.shots,
+                trainable_params=new_tape.trainable_params,
+            )
 
             num_tape += 1
 

pennylane/gradients/parameter_shift_cv.py

@@ -61,7 +61,7 @@ def _grad_method_cv(tape, idx):
             or ``"0"`` (constant parameter).
     """
 
-    par_info = tape._par_info[idx]
+    par_info = tape.par_info[idx]
     op = par_info["op"]
 
     if op.grad_method in (None, "F"):
@@ -330,7 +330,7 @@ def second_order_param_shift(tape, dev_wires, argnum=None, shifts=None, gradient
 
     for idx, _ in enumerate(tape.trainable_params):
         t_idx = list(tape.trainable_params)[idx]
-        op = tape._par_info[t_idx]["op"]
+        op = tape.par_info[t_idx]["op"]
 
         if idx not in argnum:
             # parameter has zero gradient
@@ -364,8 +364,8 @@ def second_order_param_shift(tape, dev_wires, argnum=None, shifts=None, gradient
         # evaluate transformed observables at the original parameter point
         # first build the Heisenberg picture transformation matrix Z
         Z0 = op.heisenberg_tr(dev_wires, inverse=True)
-        Z2 = shifted_tapes[0]._par_info[t_idx]["op"].heisenberg_tr(dev_wires)
-        Z1 = shifted_tapes[1]._par_info[t_idx]["op"].heisenberg_tr(dev_wires)
+        Z2 = shifted_tapes[0].par_info[t_idx]["op"].heisenberg_tr(dev_wires)
+        Z1 = shifted_tapes[1].par_info[t_idx]["op"].heisenberg_tr(dev_wires)
 
         # derivative of the operation
         Z = Z2 * coeffs[0] + Z1 * coeffs[1]
@@ -390,7 +390,7 @@ def second_order_param_shift(tape, dev_wires, argnum=None, shifts=None, gradient
 
         Z = B @ Z @ B_inv  # conjugation
 
-        g_tape = tape.copy(copy_operations=True)
+        new_measurements = list(tape.measurements)
         constants = []
 
         # transform the descendant observables into their derivatives using Z
@@ -419,9 +419,14 @@ def second_order_param_shift(tape, dev_wires, argnum=None, shifts=None, gradient
                     constant = A[0]
 
             constants.append(constant)
+            new_measurements[obs_idx] = qml.expval(op=_transform_observable(obs, Z, dev_wires))
 
-            g_tape._measurements[obs_idx] = qml.expval(op=_transform_observable(obs, Z, dev_wires))
-        g_tape._update_par_info()
+        g_tape = qml.tape.QuantumScript(
+            tape.operations,
+            new_measurements,
+            shots=tape.shots,
+            trainable_params=tape.trainable_params,
+        )
 
         if not any(i is None for i in constants):
             # Check if *all* transformed observables corresponds to a constant

pennylane/tape/operation_recorder.py

@@ -48,10 +48,9 @@ def __init__(
         ops=None,
         measurements=None,
         shots=None,
-        _update=True,
     ):  # pylint: disable=unused-argument, too-many-arguments
         AnnotatedQueue.__init__(self)
-        QuantumScript.__init__(self, ops, measurements, shots, _update=_update)
+        QuantumScript.__init__(self, ops, measurements, shots)
         self.ops = None
         self.obs = None