Fix PrepSelPrep differentiability with parameter-shift

doc/releases/changelog-dev.md

@@ -292,6 +292,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Fixes incorrect differentiation of `PrepSelPrep` when using `diff_method="parameter-shift"`. 
+  [(#6423)](https://github.com/PennyLaneAI/pennylane/pull/6423)
+
 * `default.tensor` can now handle mid circuit measurements via the deferred measurement principle.
   [(#6408)](https://github.com/PennyLaneAI/pennylane/pull/6408)
 

pennylane/templates/subroutines/prepselprep.py

@@ -38,6 +38,9 @@ def _get_new_terms(lcu):
 class PrepSelPrep(Operation):
     """Implements a block-encoding of a linear combination of unitaries.
 
+    .. warning::
+        Derivatives of this operator are not always guaranteed to exist.
+
     Args:
         lcu (Union[.Hamiltonian, .Sum, .Prod, .SProd, .LinearCombination]): The operator
             written as a linear combination of unitaries.
@@ -68,6 +71,8 @@ class PrepSelPrep(Operation):
      [ 0.75  0.25]]
     """
 
+    grad_method = None
+
     def __init__(self, lcu, control=None, id=None):
 
         coeffs, ops = lcu.terms()

tests/templates/test_subroutines/test_prepselprep.py

@@ -23,7 +23,6 @@
 import pennylane as qml
 
 
-@pytest.mark.xfail(reason="PrepSelPrep does not work with parameter-shift (GitHub issue #6331)")
 @pytest.mark.parametrize(
     ("lcu", "control"),
     [
@@ -32,10 +31,19 @@
         (qml.Hamiltonian([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]), [0]),
         (0.25 * qml.Z(2) - 0.75 * qml.X(1) @ qml.X(2), [0]),
         (qml.Z(2) + qml.X(1) @ qml.X(2), [0]),
-        (qml.ops.LinearCombination([-0.25, 0.75j], [qml.Z(3), qml.X(2) @ qml.X(3)]), [0, 1]),
-        (
+        pytest.param(
+            qml.ops.LinearCombination([-0.25, 0.75j], [qml.Z(3), qml.X(2) @ qml.X(3)]),
+            [0, 1],
+            marks=pytest.mark.skip(
+                reason="This test case is expected to fail due to `NaN`s being returned in the gradient. This is a known limitation of ``MottonenStatePreparation``."
+            ),
+        ),
+        pytest.param(
             qml.ops.LinearCombination([-0.25 + 0.1j, 0.75j], [qml.Z(4), qml.X(4) @ qml.X(5)]),
             [0, 1, 2, 3],
+            marks=pytest.mark.skip(
+                reason="This test case is expected to fail due to `NaN`s being returned in the gradient. This is a known limitation of ``MottonenStatePreparation``."
+            ),
         ),
     ],
 )

tests/templates/test_subroutines/test_qubitization.py

@@ -65,11 +65,22 @@ def circuit(theta):
     assert np.allclose(np.sort(estimated_eigenvalues), qml.eigvals(hamiltonian), atol=0.1)
 
 
-@pytest.mark.xfail(reason="PrepSelPrep does not work with parameter-shift (GitHub issue #6331)")
-def test_standard_validity():
+@pytest.mark.parametrize(
+    ("lcu", "control"),
+    [
+        (qml.dot([0.1, -0.3], [qml.X(2), qml.Z(3)]), [0]),
+        pytest.param(
+            qml.dot([0.1, -0.3, -0.3], [qml.X(0), qml.Z(1), qml.Y(0) @ qml.Z(2)]),
+            [3, 4],
+            marks=pytest.mark.skip(
+                reason="This test case is expected to fail due to `NaN`s being returned in the gradient. This is a known limitation of ``MottonenStatePreparation``."
+            ),
+        ),
+    ],
+)
+def test_standard_validity(lcu, control):
     """Check the operation using the assert_valid function."""
-    H = qml.dot([0.1, -0.3, -0.3], [qml.X(0), qml.Z(1), qml.Y(0) @ qml.Z(2)])
-    op = qml.Qubitization(H, control=[3, 4])
+    op = qml.Qubitization(lcu, control)
     qml.ops.functions.assert_valid(op)