From e9cf50c05daf3617509e16cd38ecf1847cb53dce Mon Sep 17 00:00:00 2001 From: Christina Lee Date: Mon, 27 Nov 2023 09:33:05 -0500 Subject: [PATCH] add `trainable_params` to `QuantumScript` initialization (#4877) [sc-43432] Instead of having to do: ``` tape = QuantumScript(old_tape.operations, old_tape.measurements, shots=old_tape.shots) tape.trainable_params = old_tape.trainable_params ``` We can now do: ``` tape = QuantumScript(old_tape.operations, old_tape.measurements, shots=old_tape.shots, trainable_params = old_tape.trainable_params) ``` --------- Co-authored-by: Matthew Silverman --- doc/releases/changelog-dev.md | 4 ++ pennylane/gradients/hamiltonian_grad.py | 7 ++-- pennylane/ops/functions/map_wires.py | 5 ++- pennylane/tape/qscript.py | 41 ++++++++----------- pennylane/tape/tape.py | 7 +++- pennylane/transforms/batch_input.py | 5 ++- pennylane/transforms/batch_params.py | 5 ++- pennylane/transforms/broadcast_expand.py | 5 ++- .../transforms/convert_to_numpy_parameters.py | 6 +-- pennylane/transforms/qcut/cutcircuit.py | 5 ++- tests/tape/test_qscript.py | 8 ++-- 11 files changed, 54 insertions(+), 44 deletions(-) diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index d8daaebd368..747192caaf7 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -132,6 +132,10 @@ `SparseHamiltonian`. [(#4828)](https://github.com/PennyLaneAI/pennylane/pull/4828) +* `trainable_params` can now be set on initialization of `QuantumScript`, instead of having to set the + parameter after initialization. + [(#4877)](https://github.com/PennyLaneAI/pennylane/pull/4877) + * `default.qubit` now calculates the expectation value of Hermitians in a differentiable manner. [(#4866)](https://github.com/PennyLaneAI/pennylane/pull/4866) diff --git a/pennylane/gradients/hamiltonian_grad.py b/pennylane/gradients/hamiltonian_grad.py index 7aee1997db8..211b186580d 100644 --- a/pennylane/gradients/hamiltonian_grad.py +++ b/pennylane/gradients/hamiltonian_grad.py @@ -26,14 +26,13 @@ def hamiltonian_grad(tape, idx): """ op, m_pos, p_idx = tape.get_operation(idx) - new_tape = tape.copy(copy_operations=True) # get position in queue queue_position = m_pos - len(tape.operations) - new_tape._measurements[queue_position] = qml.expval(op.ops[p_idx]) + new_measurements = list(tape.measurements) + new_measurements[queue_position] = qml.expval(op.ops[p_idx]) - new_tape._par_info = {} - new_tape._update() + new_tape = qml.tape.QuantumScript(tape.operations, new_measurements, shots=tape.shots) if len(tape.measurements) > 1: diff --git a/pennylane/ops/functions/map_wires.py b/pennylane/ops/functions/map_wires.py index 36133c50586..0a2e5ed7e94 100644 --- a/pennylane/ops/functions/map_wires.py +++ b/pennylane/ops/functions/map_wires.py @@ -112,8 +112,9 @@ def _map_wires_transform( ops = [map_wires(op, wire_map, queue=queue) for op in tape.operations] measurements = [map_wires(m, wire_map, queue=queue) for m in tape.measurements] - out = tape.__class__(ops=ops, measurements=measurements, shots=tape.shots) - out.trainable_params = tape.trainable_params + out = tape.__class__( + ops=ops, measurements=measurements, shots=tape.shots, trainable_params=tape.trainable_params + ) def processing_fn(res): """Defines how matrix works if applied to a tape containing multiple operations.""" diff --git a/pennylane/tape/qscript.py b/pennylane/tape/qscript.py index 3c3635b066a..4e3f5d9bbe6 100644 --- a/pennylane/tape/qscript.py +++ b/pennylane/tape/qscript.py @@ -93,6 +93,7 @@ class QuantumScript: Keyword Args: shots (None, int, Sequence[int], ~.Shots): Number and/or batches of shots for execution. Note that this property is still experimental and under development. + trainable_params (None, Sequence[int]): the indices for which parameters are trainable _update=True (bool): Whether or not to set various properties on initialization. Setting ``_update=False`` reduces computations if the script is only an intermediary step. @@ -176,15 +177,14 @@ def _flatten(self): @classmethod def _unflatten(cls, data, metadata): - new_tape = cls(*data, shots=metadata[0]) - new_tape.trainable_params = metadata[1] - return new_tape + return cls(*data, shots=metadata[0], trainable_params=metadata[1]) def __init__( self, ops=None, measurements=None, shots: Optional[Union[int, Sequence, Shots]] = None, + trainable_params: Optional[Sequence[int]] = None, _update=True, ): self._ops = [] if ops is None else list(ops) @@ -195,7 +195,7 @@ def __init__( """list[dict[str, Operator or int]]: Parameter information. Values are dictionaries containing the corresponding operation and operation parameter index.""" - self._trainable_params = [] + self._trainable_params = trainable_params self._graph = None self._specs = None self._output_dim = 0 @@ -433,9 +433,6 @@ def _update(self): self._update_circuit_info() # Updates wires, num_wires; O(ops+obs) self._update_par_info() # Updates _par_info; O(ops+obs) - # The following line requires _par_info to be up to date - self._update_trainable_params() # Updates the _trainable_params; O(1) - self._update_observables() # Updates _obs_sharing_wires and _obs_sharing_wires_id self._update_batch_size() # Updates _batch_size; O(ops) @@ -473,16 +470,6 @@ def _update_par_info(self): for i, d in enumerate(m.obs.data) ) - def _update_trainable_params(self): - """Set the trainable parameters - - Sets: - _trainable_params (list[int]): Script parameter indices of trainable parameters - - Call `_update_par_info` before `_update_trainable_params` - """ - self._trainable_params = list(range(len(self._par_info))) - def _update_observables(self): """Update information about observables, including the wires that are acted upon and identifying any observables that share wires. @@ -592,6 +579,8 @@ def trainable_params(self): >>> qscript.get_parameters() [0.432] """ + if self._trainable_params is None: + self._trainable_params = list(range(len(self._par_info))) return self._trainable_params @trainable_params.setter @@ -763,10 +752,12 @@ def bind_new_parameters(self, params: Sequence[TensorLike], indices: Sequence[in new_operations = new_ops[: len(self.operations)] new_measurements = new_ops[len(self.operations) :] - new_tape = self.__class__(new_operations, new_measurements, shots=self.shots) - new_tape.trainable_params = self.trainable_params - - return new_tape + return self.__class__( + new_operations, + new_measurements, + shots=self.shots, + trainable_params=self.trainable_params, + ) # ======================================================== # MEASUREMENT SHAPE @@ -886,13 +877,17 @@ def copy(self, copy_operations=False): _ops = self.operations.copy() _measurements = self.measurements.copy() - new_qscript = self.__class__(ops=_ops, measurements=_measurements, shots=self.shots) + new_qscript = self.__class__( + ops=_ops, + measurements=_measurements, + shots=self.shots, + trainable_params=list(self.trainable_params), + ) new_qscript._graph = None if copy_operations else self._graph new_qscript._specs = None new_qscript.wires = copy.copy(self.wires) new_qscript.num_wires = self.num_wires new_qscript._update_par_info() - new_qscript.trainable_params = self.trainable_params.copy() new_qscript._obs_sharing_wires = self._obs_sharing_wires new_qscript._obs_sharing_wires_id = self._obs_sharing_wires_id new_qscript._batch_size = self.batch_size diff --git a/pennylane/tape/tape.py b/pennylane/tape/tape.py index 588018eb327..7b58abe9f4c 100644 --- a/pennylane/tape/tape.py +++ b/pennylane/tape/tape.py @@ -297,6 +297,7 @@ class QuantumTape(QuantumScript, AnnotatedQueue): Keyword Args: shots (None, int, Sequence[int], ~.Shots): Number and/or batches of shots for execution. Note that this property is still experimental and under development. + trainable_params (None, Sequence[int]): the indices for which parameters are trainable _update=True (bool): Whether or not to set various properties on initialization. Setting ``_update=False`` reduces computations if the tape is only an intermediary step. @@ -420,10 +421,13 @@ def __init__( ops=None, measurements=None, shots=None, + trainable_params=None, _update=True, ): # pylint: disable=too-many-arguments AnnotatedQueue.__init__(self) - QuantumScript.__init__(self, ops, measurements, shots, _update=_update) + QuantumScript.__init__( + self, ops, measurements, shots, trainable_params=trainable_params, _update=_update + ) def __enter__(self): QuantumTape._lock.acquire() @@ -435,6 +439,7 @@ def __exit__(self, exception_type, exception_value, traceback): QueuingManager.remove_active_queue() QuantumTape._lock.release() self._process_queue() + self._trainable_params = None def adjoint(self): adjoint_tape = super().adjoint() diff --git a/pennylane/transforms/batch_input.py b/pennylane/transforms/batch_input.py index 8f0baefdad8..532c6e991da 100644 --- a/pennylane/transforms/batch_input.py +++ b/pennylane/transforms/batch_input.py @@ -101,8 +101,9 @@ def circuit(inputs, weights): output_tapes = [] for ops in _split_operations(tape.operations, all_parameters, argnum, batch_size): - new_tape = qml.tape.QuantumScript(ops, tape.measurements, shots=tape.shots) - new_tape.trainable_params = tape.trainable_params + new_tape = qml.tape.QuantumScript( + ops, tape.measurements, shots=tape.shots, trainable_params=tape.trainable_params + ) output_tapes.append(new_tape) def processing_fn(res): diff --git a/pennylane/transforms/batch_params.py b/pennylane/transforms/batch_params.py index 488a10fd8d1..1f089d0e28f 100644 --- a/pennylane/transforms/batch_params.py +++ b/pennylane/transforms/batch_params.py @@ -198,8 +198,9 @@ def circuit(x, weights): output_tapes = [] for ops in _split_operations(tape.operations, params, indices, batch_dim): - new_tape = qml.tape.QuantumScript(ops, tape.measurements, shots=tape.shots) - new_tape.trainable_params = tape.trainable_params + new_tape = qml.tape.QuantumScript( + ops, tape.measurements, shots=tape.shots, trainable_params=tape.trainable_params + ) output_tapes.append(new_tape) def processing_fn(res): diff --git a/pennylane/transforms/broadcast_expand.py b/pennylane/transforms/broadcast_expand.py index feada25e23c..43999807e6e 100644 --- a/pennylane/transforms/broadcast_expand.py +++ b/pennylane/transforms/broadcast_expand.py @@ -137,8 +137,9 @@ def null_postprocessing(results): output_tapes = [] for ops in new_ops: - new_tape = qml.tape.QuantumScript(ops, tape.measurements, shots=tape.shots) - new_tape.trainable_params = tape.trainable_params + new_tape = qml.tape.QuantumScript( + ops, tape.measurements, shots=tape.shots, trainable_params=tape.trainable_params + ) output_tapes.append(new_tape) def processing_fn(results: qml.typing.ResultBatch) -> qml.typing.Result: diff --git a/pennylane/transforms/convert_to_numpy_parameters.py b/pennylane/transforms/convert_to_numpy_parameters.py index 9908d357ba7..ed4a910b241 100644 --- a/pennylane/transforms/convert_to_numpy_parameters.py +++ b/pennylane/transforms/convert_to_numpy_parameters.py @@ -83,7 +83,7 @@ def convert_to_numpy_parameters(circuit: QuantumScript) -> QuantumScript: """ new_ops = (_convert_op_to_numpy_data(op) for op in circuit.operations) new_measurements = (_convert_measurement_to_numpy_data(m) for m in circuit.measurements) - new_circuit = circuit.__class__(new_ops, new_measurements, shots=circuit.shots) - # must preserve trainable params as we lose information about the machine learning interface - new_circuit.trainable_params = circuit.trainable_params + new_circuit = circuit.__class__( + new_ops, new_measurements, shots=circuit.shots, trainable_params=circuit.trainable_params + ) return new_circuit diff --git a/pennylane/transforms/qcut/cutcircuit.py b/pennylane/transforms/qcut/cutcircuit.py index e096a96394d..80a955927db 100644 --- a/pennylane/transforms/qcut/cutcircuit.py +++ b/pennylane/transforms/qcut/cutcircuit.py @@ -57,8 +57,9 @@ def processing_fn(res): ) new_meas_op = type(tape_meas_ops[0])(obs=qml.Hamiltonian(*tape_meas_ops[0].obs.terms())) - new_tape = qml.tape.QuantumScript(tape.operations, [new_meas_op], shots=tape.shots) - new_tape.trainable_params = tape.trainable_params + new_tape = type(tape)( + tape.operations, [new_meas_op], shots=tape.shots, trainable_params=tape.trainable_params + ) tapes, tapes_fn = qml.transforms.hamiltonian_expand(new_tape, group=False) diff --git a/tests/tape/test_qscript.py b/tests/tape/test_qscript.py index 31d96a13be1..1da39172ddb 100644 --- a/tests/tape/test_qscript.py +++ b/tests/tape/test_qscript.py @@ -34,7 +34,9 @@ def test_no_update_empty_initialization(self): assert len(qs._ops) == 0 assert len(qs._measurements) == 0 assert len(qs._par_info) == 0 - assert len(qs._trainable_params) == 0 + assert qs._trainable_params is None + assert qs.trainable_params == [] + assert qs._trainable_params == [] assert qs._graph is None assert qs._specs is None assert qs._shots.total_shots is None @@ -187,7 +189,7 @@ def test_update_par_info_update_trainable_params(self): assert p_i[6] == {"op": ops[3], "op_idx": 3, "p_idx": 1} assert p_i[7] == {"op": m[0].obs, "op_idx": 4, "p_idx": 0} - assert qs._trainable_params == list(range(8)) + assert qs.trainable_params == list(range(8)) # pylint: disable=unbalanced-tuple-unpacking def test_get_operation(self): @@ -1447,7 +1449,7 @@ def test_flatten_unflatten(qscript_type): assert all(o1 is o2 for o1, o2 in zip(new_tape.operations, tape.operations)) assert all(o1 is o2 for o1, o2 in zip(new_tape.measurements, tape.measurements)) assert new_tape.shots == qml.measurements.Shots(100) - assert new_tape.trainable_params == [0] + assert new_tape.trainable_params == (0,) @pytest.mark.jax