Support the generalization of basis state preparation and the facade legacy device for MPI LGPU (#864)

### Before submitting

Please complete the following checklist when submitting a PR:

- [x] All new features must include a unit test.
If you've fixed a bug or added code that should be tested, add a test to
the
      [`tests`](../tests) directory!

- [x] All new functions and code must be clearly commented and
documented.
If you do make documentation changes, make sure that the docs build and
      render correctly by running `make docs`.

- [x] Ensure that the test suite passes, by running `make test`.

- [x] Add a new entry to the `.github/CHANGELOG.md` file, summarizing
the
      change, and including a link back to the PR.

- [x] Ensure that code is properly formatted by running `make format`. 

When all the above are checked, delete everything above the dashed
line and fill in the pull request template.


------------------------------------------------------------------------------------------------------------

**Context:**

- PR PennyLaneAI/pennylane#6021 removed code
duplication for `BasisEmbedding` and `BasisState`. As the result
`BasisState` no longer decomposes to `BasisStatePreparation`. This PR
updates Python unit tests to support this generalization of basis state
preparation.

- PR PennyLaneAI/pennylane#6046 added a facade
wrapper class for "legacy" devices. This PR is a follow up to PR #839
updating Multi-GPU LGPU device and tests.

**Description of the Change:**

**Benefits:**

**Possible Drawbacks:**

**Related GitHub Issues:**

---------

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

Author: maliasadi

.github/CHANGELOG.md

@@ -33,11 +33,15 @@
 
 ### Improvements
 
+* Update Lightning tests to support the generalization of basis state preparation.
+  [(#864)](https://github.com/PennyLaneAI/pennylane-lightning/pull/864)
+
 * 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)
+  [(#864)](https://github.com/PennyLaneAI/pennylane-lightning/pull/864)
 
 * LightningQubit gains native support for the `PauliRot` gate.
   [(#834)](https://github.com/PennyLaneAI/pennylane-lightning/pull/834)

mpitests/test_adjoint_jacobian.py

@@ -649,25 +649,20 @@ def dev(self, request):
             batch_obs=request.param[1],
         )
 
-    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=8, mpi=True, 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):
@@ -689,7 +684,7 @@ def circuit(x, y, z):
             return qml.expval(qml.PauliX(0) @ qml.PauliZ(1))
 
         qnode1 = QNode(circuit, dev, diff_method="adjoint")
-        spy = mocker.spy(dev, "adjoint_jacobian")
+        spy = mocker.spy(dev.target_device, "adjoint_jacobian")
 
         grad_fn = qml.grad(qnode1)
         grad_A = grad_fn(*args)
@@ -731,7 +726,7 @@ def cost(p1, p2):
         zero_state = np.array([1.0, 0.0])
         cost(reused_p, other_p)
 
-        spy = mocker.spy(dev, "adjoint_jacobian")
+        spy = mocker.spy(dev.target_device, "adjoint_jacobian")
 
         # analytic gradient
         grad_fn = qml.grad(cost)
@@ -770,7 +765,7 @@ def circuit(params):
             qml.Rot(params[1], params[0], 2 * params[0], wires=[0])
             return qml.expval(qml.PauliX(0))
 
-        spy_analytic = mocker.spy(dev, "adjoint_jacobian")
+        spy_analytic = mocker.spy(dev.target_device, "adjoint_jacobian")
 
         h = 1e-3 if dev.R_DTYPE == np.float32 else 1e-7
         tol = 1e-3 if dev.R_DTYPE == np.float32 else 1e-7

mpitests/test_apply.py

@@ -672,29 +672,30 @@ def test_sample_dimensions(self, C_DTYPE):
         """
         num_wires = numQubits
 
-        dev = qml.device("lightning.gpu", wires=num_wires, mpi=True, shots=1000, c_dtype=C_DTYPE)
-
-        dev.apply([qml.RX(1.5708, wires=[0]), qml.RX(1.5708, wires=[1])])
-
-        dev.shots = 10
-        dev._wires_measured = {0}
-        dev._samples = dev.generate_samples()
-        s1 = dev.sample(qml.PauliZ(wires=[0]))
-        assert np.array_equal(s1.shape, (10,))
-
-        dev.reset()
-        dev.shots = 12
-        dev._wires_measured = {1}
-        dev._samples = dev.generate_samples()
-        s2 = dev.sample(qml.PauliZ(wires=[1]))
-        assert np.array_equal(s2.shape, (12,))
-
-        dev.reset()
-        dev.shots = 17
-        dev._wires_measured = {0, 1}
-        dev._samples = dev.generate_samples()
-        s3 = dev.sample(qml.PauliX(0) @ qml.PauliZ(1))
-        assert np.array_equal(s3.shape, (17,))
+        dev = qml.device("lightning.gpu", wires=num_wires, mpi=True, c_dtype=C_DTYPE)
+
+        ops = [qml.RX(1.5708, wires=[0]), qml.RX(1.5708, wires=[1])]
+
+        shots = 10
+        obs = qml.PauliZ(wires=[0])
+        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])
+        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)
+        tape = qml.tape.QuantumScript(ops, [qml.sample(op=obs)], shots=shots)
+        s3 = dev.execute(tape)
+
+        assert np.array_equal(s3.shape, (shots,))
 
     @pytest.mark.parametrize("C_DTYPE", [np.complex128, np.complex64])
     def test_sample_values(self, tol, C_DTYPE):
@@ -703,13 +704,13 @@ def test_sample_values(self, tol, C_DTYPE):
         """
         num_wires = numQubits
 
-        dev = qml.device("lightning.gpu", wires=num_wires, mpi=True, shots=1000, c_dtype=C_DTYPE)
-        dev.reset()
-        dev.apply([qml.RX(1.5708, wires=[0])])
-        dev._wires_measured = {0}
-        dev._samples = dev.generate_samples()
+        dev = qml.device("lightning.gpu", wires=num_wires, mpi=True, c_dtype=C_DTYPE)
 
-        s1 = dev.sample(qml.PauliZ(0))
+        shots = qml.measurements.Shots(1000)
+        ops = [qml.RX(1.5708, wires=[0])]
+        obs = qml.PauliZ(0)
+        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

pennylane_lightning/core/_version.py

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

tests/lightning_qubit/test_state_vector_class.py

@@ -58,10 +58,10 @@ def test_wrong_dtype(dtype):
 
 
 def test_errors_basis_state():
-    with pytest.raises(ValueError, match="BasisState parameter must consist of 0 or 1 integers."):
+    with pytest.raises(ValueError, match="Basis state must only consist of 0s and 1s;"):
         state_vector = LightningStateVector(2)
         state_vector.apply_operations([qml.BasisState(np.array([-0.2, 4.2]), wires=[0, 1])])
-    with pytest.raises(ValueError, match="BasisState parameter and wires must be of equal length."):
+    with pytest.raises(ValueError, match="State must be of length 1;"):
         state_vector = LightningStateVector(1)
         state_vector.apply_operations([qml.BasisState(np.array([0, 1]), wires=[0])])
 

tests/lightning_tensor/test_tensornet_class.py

@@ -57,10 +57,10 @@ def test_wrong_device_name():
 
 def test_errors_basis_state():
     """Test that errors are raised when applying a BasisState operation."""
-    with pytest.raises(ValueError, match="BasisState parameter must consist of 0 or 1 integers."):
+    with pytest.raises(ValueError, match="Basis state must only consist of 0s and 1s;"):
         tensornet = LightningTensorNet(3, 5)
         tensornet.apply_operations([qml.BasisState(np.array([-0.2, 4.2]), wires=[0, 1])])
-    with pytest.raises(ValueError, match="BasisState parameter and wires must be of equal length."):
+    with pytest.raises(ValueError, match="State must be of length 1;"):
         tensornet = LightningTensorNet(3, 5)
         tensornet.apply_operations([qml.BasisState(np.array([0, 1]), wires=[0])])
 

tests/new_api/test_device.py

@@ -349,8 +349,9 @@ def test_preprocess_state_prep_first_op_decomposition(self, op, is_trainable):
         device = LightningDevice(wires=3)
 
         if is_trainable:
-            # Need to decompose twice as the state prep ops we use first decompose into a template
-            decomp = op.decomposition()[0].decomposition()
+            decomp = op.decomposition()
+            # decompose one more time if it's decomposed into a template:
+            decomp = decomp[0].decomposition() if len(decomp) == 1 else decomp
         else:
             decomp = [op]
 
@@ -367,7 +368,7 @@ def test_preprocess_state_prep_first_op_decomposition(self, op, is_trainable):
             (qml.StatePrep(np.array([1, 0]), wires=0), 1),
             (qml.BasisState([1, 1], wires=[0, 1]), 1),
             (qml.BasisState(qml.numpy.array([1, 1]), wires=[0, 1]), 1),
-            (qml.AmplitudeEmbedding([1 / np.sqrt(2), 1 / np.sqrt(2)], wires=0), 2),
+            (qml.AmplitudeEmbedding([1 / np.sqrt(2), 1 / np.sqrt(2)], wires=0), 1),
             (qml.MottonenStatePreparation([1 / np.sqrt(2), 1 / np.sqrt(2)], wires=0), 0),
         ],
     )
@@ -378,8 +379,7 @@ def test_preprocess_state_prep_middle_op_decomposition(self, op, decomp_depth):
         )
         device = LightningDevice(wires=3)
 
-        for _ in range(decomp_depth):
-            op = op.decomposition()[0]
+        op = op.decomposition()[0] if decomp_depth and len(op.decomposition()) == 1 else op
         decomp = op.decomposition()
 
         program, _ = device.preprocess()

tests/test_apply.py

@@ -488,7 +488,7 @@ def test_apply_operation_preserve_pointer_two_wires_with_parameters(
     def test_apply_errors_qubit_state_vector(self, stateprep, qubit_device):
         """Test that apply fails for incorrect state preparation, and > 2 qubit gates"""
         dev = qubit_device(wires=2)
-        with pytest.raises(ValueError, match="Sum of amplitudes-squared does not equal one."):
+        with pytest.raises(ValueError, match="The state must be a vector of norm 1.0;"):
             dev.apply([stateprep(np.array([1, -1]), wires=[0])])
 
         with pytest.raises(
@@ -500,14 +500,10 @@ def test_apply_errors_qubit_state_vector(self, stateprep, qubit_device):
 
     def test_apply_errors_basis_state(self, qubit_device):
         dev = qubit_device(wires=2)
-        with pytest.raises(
-            ValueError, match="BasisState parameter must consist of 0 or 1 integers."
-        ):
+        with pytest.raises(ValueError, match="Basis state must only consist of 0s and 1s;"):
             dev.apply([qml.BasisState(np.array([-0.2, 4.2]), wires=[0, 1])])
 
-        with pytest.raises(
-            ValueError, match="BasisState parameter and wires must be of equal length."
-        ):
+        with pytest.raises(ValueError, match="State must be of length 1;"):
             dev.apply([qml.BasisState(np.array([0, 1]), wires=[0])])
 
         with pytest.raises(