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 <[email protected]>
PennyLaneAI · Nov 27, 2023 · e9cf50c · e9cf50c
1 parent d6a8049
commit e9cf50c
Show file tree

Hide file tree

Showing 11 changed files with 54 additions and 44 deletions.
diff --git 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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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
@@ -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