Skip to content
New issue

Have a question about this project? Sign up for a free GitHub account to open an issue and contact its maintainers and the community.

By clicking “Sign up for GitHub”, you agree to our terms of service and privacy statement. We’ll occasionally send you account related emails.

Already on GitHub? Sign in to your account

Lazy loading of Quantumscript properties #5696

Merged
merged 14 commits into from
May 24, 2024
Merged
Show file tree
Hide file tree
Changes from all commits
Commits
File filter

Filter by extension

Filter by extension

Conversations
Failed to load comments.
Loading
Jump to
Jump to file
Failed to load files.
Loading
Diff view
Diff view
5 changes: 5 additions & 0 deletions doc/releases/changelog-dev.md
Original file line number Diff line number Diff line change
Expand Up @@ -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`.
Expand Down Expand Up @@ -197,6 +201,7 @@ Astral Cai,
Ahmed Darwish,
Isaac De Vlugt,
Pietropaolo Frisoni,
Emiliano Godinez,
Soran Jahangiri,
Korbinian Kottmann,
Christina Lee,
Expand Down
14 changes: 5 additions & 9 deletions pennylane/_device.py
Original file line number Diff line number Diff line change
Expand Up @@ -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
Expand Down Expand Up @@ -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
Expand All @@ -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
Expand Down Expand Up @@ -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)

Expand All @@ -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

Expand Down Expand Up @@ -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
Expand Down
2 changes: 1 addition & 1 deletion pennylane/_qubit_device.py
Original file line number Diff line number Diff line change
Expand Up @@ -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 "
Expand Down
2 changes: 1 addition & 1 deletion pennylane/fourier/qnode_spectrum.py
Original file line number Diff line number Diff line change
Expand Up @@ -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):
Expand Down
2 changes: 1 addition & 1 deletion pennylane/gradients/adjoint_metric_tensor.py
Original file line number Diff line number Diff line change
Expand Up @@ -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))
Expand Down
2 changes: 1 addition & 1 deletion pennylane/gradients/gradient_transform.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"]]
Expand Down
12 changes: 7 additions & 5 deletions pennylane/gradients/hadamard_gradient.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
19 changes: 12 additions & 7 deletions pennylane/gradients/parameter_shift_cv.py
Original file line number Diff line number Diff line change
Expand Up @@ -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"):
Expand Down Expand Up @@ -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
Expand Down Expand Up @@ -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]
Expand All @@ -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
Expand Down Expand Up @@ -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
Expand Down
3 changes: 1 addition & 2 deletions pennylane/tape/operation_recorder.py
Original file line number Diff line number Diff line change
Expand Up @@ -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

Expand Down
Loading
Loading