update to support None and add tests

pennylane/tape/qscript.py

@@ -375,22 +375,23 @@ def op_wires(self) -> Wires:
 
     def update(
         self,
-        operations: Optional[Iterable[Operator]] = None,
-        measurements: Optional[Iterable[MeasurementProcess]] = None,
-        shots: Optional[ShotsLike] = None,
-        trainable_params: Optional[Sequence[int]] = None,
+        **kwargs,
+        # operations: Optional[Iterable[Operator]] = "unset",
+        # measurements: Optional[Iterable[MeasurementProcess]] = "unset",
+        # shots: Optional[ShotsLike] = "unset",
+        # trainable_params: Optional[Sequence[int]] = "unset",
     ):
-        """Update attirbutes on the tape.
+        r"""update(operations, measurements, shots, trainable_params)
+
+        Update attributes on the tape.
 
         Keyword Args:
             operations (Iterable[Operator]): An iterable of the operations to replace on the tape.
-                Defaults to None.
-            measurements (Iterable[MeasurementProcess]): All the measurements to replace on the
-                tape. Defaults to None.
+            measurements (Iterable[MeasurementProcess]): An iterable of all the measurements to
+                replace on the tape.
             shots (None, int, Sequence[int], ~.Shots): Number and/or batches of shots to replace
-                on the tape. Defaults to None.
+                on the tape.
             trainable_params (None, Sequence[int]): the indices for which parameters are trainable.
-                Defaults to None.
 
         Returns: A new tape instance, initialized with any kwargs passed to update.
             Anything not set in update will be the same as the initial tape.
@@ -417,11 +418,20 @@ def update(
         >>> new_tape.shots
         2000
         """
+        for k in kwargs:
+            if k not in ["operations", "measurements", "shots", "trainable_params"]:
+                raise TypeError(
+                    f"{self.__class__}.update() got an unexpected keyword argument '{k}'"
+                )
 
-        ops = operations or self.operations
-        measurements = measurements or self.measurements
-        shots = shots or self.shots
-        trainable_params = trainable_params or self.trainable_params
+        ops = kwargs.get("operations") if "operations" in kwargs else self.operations
+        measurements = kwargs.get("measurements") if "measurements" in kwargs else self.measurements
+        shots = kwargs.get("shots") if "shots" in kwargs else self.shots
+        trainable_params = (
+            kwargs.get("trainable_params")
+            if "trainable_params" in kwargs
+            else self.trainable_params
+        )
 
         return self.__class__(ops, measurements, shots, trainable_params)
 

tests/tape/test_qscript.py

@@ -1416,52 +1416,87 @@ def test_jax_pytree_integration(qscript_type):
     assert qml.math.allclose(data[5], eye_mat)
 
 
-@pytest.mark.parametrize(
-    "update_input",
-    [
-        {"measurements": [qml.expval(qml.X(0))]},
-        {"operations": [qml.X(7)]},
-        {"shots": 100},
-        {"shots": None},
-        {"shots": 50, "measurements": [qml.sample(wires=[2, 3])]},
-        {"operations": [qml.RX(1.2, 0), qml.RY(2.3, 1)], "trainable_params": [0, 1]},
-    ],
-)
-@pytest.mark.parametrize("tape_class", [qml.tape.QuantumTape, QuantumScript])
-def test_public_update(self, update_input, tape_class):
-    """Test the public update method behaves as expected"""
-
-    initial_args = [[qml.X("b"), qml.RX(1.2, "a")]]
-    initial_kwargs = {
-        "measurements": [qml.counts()],
-        "shots": 2500,
-        "trainable_params": [1],
-    }
-
-    tape = tape_class(*initial_args, **initial_kwargs)
-
-    new_tape = tape.update(**update_input)
-
-    # original tape is unmodified
-    assert tape.operations == initial_args[0]
-    for kwarg, val in initial_kwargs.items():
-        if kwarg == "shots":
-            val = Shots(val)
-        assert getattr(tape, kwarg) == val
-
-    if "operations" in update_input:
-        assert new_tape.operations == update_input["operations"]
-        del update_input["operations"]
-    else:
-        assert new_tape.operations == tape.operations
-
-    for kwarg, val in update_input.items():
-        if kwarg == "shots":
-            val = Shots(val)
-        assert getattr(new_tape, kwarg) == val
-        del initial_kwargs[kwarg]
-
-    for kwarg, val in initial_kwargs.items():
-        if kwarg == "shots":
-            val = Shots(val)
-        assert getattr(new_tape, kwarg) == getattr(tape, kwarg)
+class TestPublicUpdate:
+    """Test the public update method"""
+
+    @pytest.mark.parametrize("shots", [50, (1000, 2000), None])
+    def test_public_update_shots(self, shots):
+        """Test the public update method behaves as expected for setting shots"""
+
+        ops = [qml.X("b"), qml.RX(1.2, "a")]
+        tape = QuantumScript(ops, measurements=[qml.counts()], shots=2500, trainable_params=[1])
+
+        new_tape = tape.update(shots=shots)
+        assert tape.shots == Shots(2500)
+        assert new_tape.shots == Shots(shots)
+
+        assert new_tape.operations == tape.operations == ops
+        assert new_tape.measurements == tape.measurements == [qml.counts()]
+        assert new_tape.trainable_params == tape.trainable_params == [1]
+
+    def test_public_update_measurements(self):
+        """Test the public update method behaves as expected for setting measurements"""
+
+        ops = [qml.X("b"), qml.RX(1.2, "a")]
+        tape = QuantumScript(ops, measurements=[qml.counts()], shots=2500, trainable_params=[1])
+
+        new_measurements = [qml.expval(qml.X(0)), qml.sample()]
+        new_tape = tape.update(measurements=new_measurements)
+
+        assert tape.measurements == [qml.counts()]
+        assert new_tape.measurements == new_measurements
+
+        assert new_tape.operations == tape.operations == ops
+        assert new_tape.shots == tape.shots == Shots(2500)
+        assert new_tape.trainable_params == tape.trainable_params == [1]
+
+    def test_public_update_operations(self):
+        """Test the public update method behaves as expected for setting operations"""
+
+        ops = [qml.X("b"), qml.RX(1.2, "a")]
+        tape = QuantumScript(ops, measurements=[qml.counts()], shots=2500, trainable_params=[1])
+
+        new_ops = [qml.X(0)]
+        new_tape = tape.update(operations=new_ops)
+
+        assert tape.operations == ops
+        assert new_tape.operations == new_ops
+
+        assert new_tape.measurements == tape.measurements == [qml.counts()]
+        assert new_tape.shots == tape.shots == Shots(2500)
+        assert new_tape.trainable_params == tape.trainable_params == [1]
+
+    def test_public_update_trainable_params(self):
+        """Test the public update method behaves as expected for setting trainable parameters"""
+
+        ops = [qml.RX(1.23, "b"), qml.RX(4.56, "a")]
+        tape = QuantumScript(ops, measurements=[qml.counts()], shots=2500, trainable_params=[1])
+
+        new_tape = tape.update(trainable_params=[0])
+
+        assert tape.trainable_params == [1]
+        assert tape.get_parameters() == [4.56]
+        assert new_tape.trainable_params == [0]
+        assert new_tape.get_parameters() == [1.23]
+
+        assert new_tape.operations == tape.operations == ops
+        assert new_tape.measurements == tape.measurements == [qml.counts()]
+        assert new_tape.shots == tape.shots == Shots(2500)
+
+    def test_public_update_bad_kwarg(self):
+        """Test that an unrecognized keyword argument raises an error"""
+
+        tape = QuantumScript([qml.X(0)], [qml.counts()], shots=2500)
+
+        with pytest.raises(TypeError, match="unexpected keyword argument"):
+            _ = tape.update(bad_kwarg=3)
+
+    # pylint: disable = unidiomatic-typecheck
+    @pytest.mark.parametrize("qscript_type", (QuantumScript, qml.tape.QuantumTape))
+    def test_public_update_preserves_class(self, qscript_type):
+        """Test that the type of the updated tape is unaltered"""
+
+        tape = qscript_type([qml.X(0)], [qml.counts()], shots=2500)
+        new_tape = tape.update(operations=[qml.Y(1)])
+
+        assert type(new_tape) == type(tape) == qscript_type