Update generate_samples in LK and LGPU to support qml.measurements.Shots

.github/CHANGELOG.md

@@ -27,6 +27,9 @@
 
 ### Improvements
 
+* 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)
 
@@ -136,7 +139,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-dev38"
+__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 @@
             ``(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])