Update generate_samples in LK and LGPU to support qml.measurements.Shots (#839)

**Context:**
PR PennyLaneAI/pennylane#6046 wraps the legacy
device API automatically in various device creation, qnode, and execute
functions. As LK and LGPU plugins still rely on the legacy device API,
the shots tests and the `generate_samples` logic in
`lightning_kokkos.py` and `lightning_gpu.py` should be updated to adhere
the new convention.

**Related Shortcut Stories:**
[sc-65998]

---------

Co-authored-by: ringo-but-quantum <[email protected]>
Co-authored-by: Shiro-Raven <[email protected]>
Co-authored-by: albi3ro <[email protected]>

Author: 4 people

.github/CHANGELOG.md

@@ -32,6 +32,9 @@
 
 * Multiple calls to the `append_mps_final_state()` API is allowed in `lightning.tensor`.
   [(#830)](https://github.com/PennyLaneAI/pennylane-lightning/pull/830)
+  
+* Update `generate_samples` in `LightningKokkos` and `LightningGPU` to support `qml.measurements.Shots` type instances.
+  [(#839)](https://github.com/PennyLaneAI/pennylane-lightning/pull/839)
 
 * LightningQubit gains native support for the `PauliRot` gate.
   [(#834)](https://github.com/PennyLaneAI/pennylane-lightning/pull/834)
@@ -142,7 +145,7 @@
 
 This release contains contributions from (in alphabetical order):
 
-Ali Asadi, Astral Cai, Amintor Dusko, Vincent Michaud-Rioux, Erick Ochoa Lopez, Lee J. O'Riordan, Mudit Pandey, Shuli Shu, Raul Torres, Paul Haochen Wang
+Ali Asadi, Astral Cai, Ahmed Darwish, Amintor Dusko, Vincent Michaud-Rioux, Erick Ochoa Lopez, Lee J. O'Riordan, Mudit Pandey, Shuli Shu, Raul Torres, Paul Haochen Wang
 
 ---
 

pennylane_lightning/core/_version.py

@@ -16,4 +16,4 @@
    Version number (major.minor.patch[-label])
 """
 
-__version__ = "0.38.0-dev40"
+__version__ = "0.38.0-dev39"

pennylane_lightning/lightning_gpu/lightning_gpu.py

@@ -38,7 +38,6 @@
 from pennylane_lightning.core.lightning_base import LightningBase
 
 try:
-
     from pennylane_lightning.lightning_gpu_ops import (
         DevPool,
         MeasurementsC64,
@@ -818,9 +817,9 @@ def generate_samples(self):
             array[int]: array of samples in binary representation with shape
             ``(dev.shots, dev.num_wires)``
         """
-        return self.measurements.generate_samples(len(self.wires), self.shots).astype(
-            int, copy=False
-        )
+        shots = self.shots if isinstance(self.shots, int) else self.shots.total_shots
+
+        return self.measurements.generate_samples(len(self.wires), shots).astype(int, copy=False)
 
     # pylint: disable=protected-access
     def expval(self, observable, shot_range=None, bin_size=None):

pennylane_lightning/lightning_kokkos/lightning_kokkos.py

@@ -614,6 +614,9 @@ def generate_samples(self, shots=None):
             ``(dev.shots, dev.num_wires)``
         """
         shots = self.shots if shots is None else shots
+
+        shots = shots.total_shots if isinstance(shots, qml.measurements.Shots) else shots
+
         measure = (
             MeasurementsC64(self._kokkos_state)
             if self.use_csingle

tests/lightning_qubit/test_measurements_samples_MCMC.py

@@ -46,15 +46,8 @@ def test_mcmc_sample_dimensions(self, dev, num_shots, measured_wires, operation,
         the correct dimensions
         """
         ops = [qml.RX(1.5708, wires=[0]), qml.RX(1.5708, wires=[1])]
-        if ld._new_API:
-            tape = qml.tape.QuantumScript(ops, [qml.sample(op=operation)], shots=num_shots)
-            s1 = dev.execute(tape)
-        else:
-            dev.apply(ops)
-            dev.shots = num_shots
-            dev._wires_measured = measured_wires
-            dev._samples = dev.generate_samples()
-            s1 = dev.sample(operation)
+        tape = qml.tape.QuantumScript(ops, [qml.sample(op=operation)], shots=num_shots)
+        s1 = dev.execute(tape)
 
         assert np.array_equal(s1.shape, (shape,))
 
@@ -67,14 +60,8 @@ def test_sample_values(self, tol, kernel):
             device_name, wires=2, shots=1000, mcmc=True, kernel_name=kernel, num_burnin=100
         )
         ops = [qml.RX(1.5708, wires=[0])]
-        if ld._new_API:
-            tape = qml.tape.QuantumScript(ops, [qml.sample(op=qml.PauliZ(0))], shots=1000)
-            s1 = dev.execute(tape)
-        else:
-            dev.apply([qml.RX(1.5708, wires=[0])])
-            dev._wires_measured = {0}
-            dev._samples = dev.generate_samples()
-            s1 = dev.sample(qml.PauliZ(0))
+        tape = qml.tape.QuantumScript(ops, [qml.sample(op=qml.PauliZ(0))], shots=1000)
+        s1 = dev.execute(tape)
 
         # s1 should only contain 1 and -1, which is guaranteed if
         # they square to 1

tests/lightning_tensor/test_tensornet_class.py

@@ -21,7 +21,6 @@
 import pennylane as qml
 import pytest
 from conftest import LightningDevice, device_name  # tested device
-from pennylane import DeviceError
 from pennylane.wires import Wires
 
 if device_name != "lightning.tensor":
@@ -88,6 +87,7 @@ def test_errors_apply_operation_state_preparation(operation, par):
     tensornet = LightningTensorNet(wires, bondDims)
 
     with pytest.raises(
-        DeviceError, match="lightning.tensor does not support initialization with a state vector."
+        qml.DeviceError,
+        match="lightning.tensor does not support initialization with a state vector.",
     ):
         tensornet.apply_operations([operation(np.array(par), Wires(range(wires)))])

tests/test_adjoint_jacobian.py

@@ -700,23 +700,20 @@ def dev(self, request):
         return qml.device(device_name, wires=2, c_dtype=request.param)
 
     @pytest.mark.skipif(ld._new_API, reason="Old API required")
-    def test_finite_shots_warning(self):
-        """Tests that a warning is raised when computing the adjoint diff on a device with finite shots"""
+    def test_finite_shots_error(self):
+        """Tests that an error is raised when computing the adjoint diff on a device with finite shots"""
 
         dev = qml.device(device_name, wires=1, shots=1)
 
-        with pytest.warns(
-            UserWarning, match="Requested adjoint differentiation to be computed with finite shots."
+        with pytest.raises(
+            qml.QuantumFunctionError, match="does not support adjoint with requested circuit."
         ):
 
             @qml.qnode(dev, diff_method="adjoint")
             def circ(x):
                 qml.RX(x, wires=0)
                 return qml.expval(qml.PauliZ(0))
 
-        with pytest.warns(
-            UserWarning, match="Requested adjoint differentiation to be computed with finite shots."
-        ):
             qml.grad(circ)(0.1)
 
     def test_qnode(self, mocker, dev):
@@ -741,7 +738,7 @@ def circuit(x, y, z):
         spy = (
             mocker.spy(dev, "execute_and_compute_derivatives")
             if ld._new_API
-            else mocker.spy(dev, "adjoint_jacobian")
+            else mocker.spy(dev.target_device, "adjoint_jacobian")
         )
         tol, h = get_tolerance_and_stepsize(dev, step_size=True)
 
@@ -926,7 +923,7 @@ def cost(p1, p2):
         if ld._new_API:
             spy = mocker.spy(dev, "execute_and_compute_derivatives")
         else:
-            spy = mocker.spy(dev, "adjoint_jacobian")
+            spy = mocker.spy(dev.target_device, "adjoint_jacobian")
 
         # analytic gradient
         grad_fn = qml.grad(cost)
@@ -968,7 +965,7 @@ def circuit(params):
         spy_analytic = (
             mocker.spy(dev, "execute_and_compute_derivatives")
             if ld._new_API
-            else mocker.spy(dev, "adjoint_jacobian")
+            else mocker.spy(dev.target_device, "adjoint_jacobian")
         )
         tol, h = get_tolerance_and_stepsize(dev, step_size=True)
 

tests/test_apply.py

@@ -566,13 +566,8 @@ def test_expval_single_wire_no_parameters(
         dev = qubit_device(wires=1)
         obs = operation(wires=[0])
         ops = [stateprep(np.array(input), wires=[0])]
-        if ld._new_API:
-            tape = qml.tape.QuantumScript(ops, [qml.expval(op=obs)])
-            res = dev.execute(tape)
-        else:
-            dev.reset()
-            dev.apply(ops, obs.diagonalizing_gates())
-            res = dev.expval(obs)
+        tape = qml.tape.QuantumScript(ops, [qml.expval(op=obs)])
+        res = dev.execute(tape)
 
         assert np.isclose(res, expected_output, atol=tol, rtol=0)
 
@@ -630,13 +625,8 @@ def test_var_single_wire_no_parameters(
         dev = qubit_device(wires=1)
         obs = operation(wires=[0])
         ops = [stateprep(np.array(input), wires=[0])]
-        if ld._new_API:
-            tape = qml.tape.QuantumScript(ops, [qml.var(op=obs)])
-            res = dev.execute(tape)
-        else:
-            dev.reset()
-            dev.apply(ops, obs.diagonalizing_gates())
-            res = dev.var(obs)
+        tape = qml.tape.QuantumScript(ops, [qml.var(op=obs)])
+        res = dev.execute(tape)
 
         assert np.isclose(res, expected_output, atol=tol, rtol=0)
 
@@ -680,42 +670,22 @@ def test_sample_dimensions(self, qubit_device):
 
         shots = 10
         obs = qml.PauliZ(wires=[0])
-        if ld._new_API:
-            tape = qml.tape.QuantumScript(ops, [qml.sample(op=obs)], shots=shots)
-            s1 = dev.execute(tape)
-        else:
-            dev.reset()
-            dev.apply(ops)
-            dev.shots = shots
-            dev._wires_measured = {0}
-            dev._samples = dev.generate_samples()
-            s1 = dev.sample(obs)
+        tape = qml.tape.QuantumScript(ops, [qml.sample(op=obs)], shots=shots)
+        s1 = dev.execute(tape)
+
         assert np.array_equal(s1.shape, (shots,))
 
         shots = 12
         obs = qml.PauliZ(wires=[1])
-        if ld._new_API:
-            tape = qml.tape.QuantumScript(ops, [qml.sample(op=obs)], shots=shots)
-            s2 = dev.execute(tape)
-        else:
-            dev.reset()
-            dev.shots = shots
-            dev._wires_measured = {1}
-            dev._samples = dev.generate_samples()
-            s2 = dev.sample(qml.PauliZ(wires=[1]))
+        tape = qml.tape.QuantumScript(ops, [qml.sample(op=obs)], shots=shots)
+        s2 = dev.execute(tape)
         assert np.array_equal(s2.shape, (shots,))
 
         shots = 17
         obs = qml.PauliX(0) @ qml.PauliZ(1)
-        if ld._new_API:
-            tape = qml.tape.QuantumScript(ops, [qml.sample(op=obs)], shots=shots)
-            s3 = dev.execute(tape)
-        else:
-            dev.reset()
-            dev.shots = shots
-            dev._wires_measured = {0, 1}
-            dev._samples = dev.generate_samples()
-            s3 = dev.sample(qml.PauliZ(wires=[1]))
+        tape = qml.tape.QuantumScript(ops, [qml.sample(op=obs)], shots=shots)
+        s3 = dev.execute(tape)
+
         assert np.array_equal(s3.shape, (shots,))
 
     def test_sample_values(self, qubit_device, tol):
@@ -730,18 +700,10 @@ def test_sample_values(self, qubit_device, tol):
 
         ops = [qml.RX(1.5708, wires=[0])]
 
-        shots = 1000
+        shots = qml.measurements.Shots(1000)
         obs = qml.PauliZ(0)
-        if ld._new_API:
-            tape = qml.tape.QuantumScript(ops, [qml.sample(op=obs)], shots=shots)
-            s1 = dev.execute(tape)
-        else:
-            dev.reset()
-            dev.apply(ops)
-            dev.shots = shots
-            dev._wires_measured = {0}
-            dev._samples = dev.generate_samples()
-            s1 = dev.sample(obs)
+        tape = qml.tape.QuantumScript(ops, [qml.sample(op=obs)], shots=shots)
+        s1 = dev.execute(tape)
 
         # s1 should only contain 1 and -1, which is guaranteed if
         # they square to 1
@@ -756,13 +718,8 @@ def test_load_default_qubit_device(self):
         """Test that the default plugin loads correctly"""
 
         dev = qml.device(device_name, wires=2)
-        if dev._new_API:
-            assert not dev.shots
-            assert len(dev.wires) == 2
-        else:
-            assert dev.shots is None
-            assert dev.num_wires == 2
-            assert dev.short_name == device_name
+        assert not dev.shots
+        assert len(dev.wires) == 2
 
     @pytest.mark.xfail(ld._new_API, reason="Old device API required.")
     def test_no_backprop(self):
@@ -1276,14 +1233,10 @@ def test_multi_samples_return_correlated_results(self, qubit_device):
         def circuit():
             qml.Hadamard(0)
             qml.CNOT(wires=[0, 1])
-            if ld._new_API:
-                return qml.sample(wires=[0, 1])
-            else:
-                return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1))
+            return qml.sample(wires=[0, 1])
 
         outcomes = circuit()
-        if ld._new_API:
-            outcomes = outcomes.T
+        outcomes = outcomes.T
 
         assert np.array_equal(outcomes[0], outcomes[1])
 
@@ -1305,14 +1258,10 @@ def test_multi_samples_return_correlated_results_more_wires_than_size_of_observa
         def circuit():
             qml.Hadamard(0)
             qml.CNOT(wires=[0, 1])
-            if ld._new_API:
-                return qml.sample(wires=[0, 1])
-            else:
-                return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1))
+            return qml.sample(wires=[0, 1])
 
         outcomes = circuit()
-        if ld._new_API:
-            outcomes = outcomes.T
+        outcomes = outcomes.T
 
         assert np.array_equal(outcomes[0], outcomes[1])
 
@@ -1350,14 +1299,10 @@ def circuit():
             qml.Snapshot()
             qml.adjoint(qml.Snapshot())
             qml.CNOT(wires=[0, 1])
-            if ld._new_API:
-                return qml.sample(wires=[0, 1])
-            else:
-                return qml.sample(qml.PauliZ(0)), qml.sample(qml.PauliZ(1))
+            return qml.sample(wires=[0, 1])
 
         outcomes = circuit()
-        if ld._new_API:
-            outcomes = outcomes.T
+        outcomes = outcomes.T
 
         assert np.array_equal(outcomes[0], outcomes[1])