Fix PrepSelPrep differentiability with parameter-shift

doc/releases/changelog-dev.md

@@ -312,6 +312,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,27 +23,28 @@
 import pennylane as qml
 
 
-@pytest.mark.xfail(reason="PrepSelPrep does not work with parameter-shift (GitHub issue #6331)")
 @pytest.mark.parametrize(
-    ("lcu", "control"),
+    ("lcu", "control", "skip_diff"),
     [
-        (qml.ops.LinearCombination([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]), [0]),
-        (qml.dot([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]), [0]),
-        (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]),
+        (qml.ops.LinearCombination([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]), [0], False),
+        (qml.dot([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]), [0], False),
+        (qml.Hamiltonian([0.25, 0.75], [qml.Z(2), qml.X(1) @ qml.X(2)]), [0], False),
+        (0.25 * qml.Z(2) - 0.75 * qml.X(1) @ qml.X(2), [0], False),
+        (qml.Z(2) + qml.X(1) @ qml.X(2), [0], False),
+        (qml.ops.LinearCombination([-0.25, 0.75j], [qml.Z(3), qml.X(2) @ qml.X(3)]), [0, 1], True),
         (
             qml.ops.LinearCombination([-0.25 + 0.1j, 0.75j], [qml.Z(4), qml.X(4) @ qml.X(5)]),
             [0, 1, 2, 3],
+            True,
         ),
     ],
 )
-def test_standard_checks(lcu, control):
+def test_standard_checks(lcu, control, skip_diff):
     """Run standard validity tests."""
 
     op = qml.PrepSelPrep(lcu, control)
-    qml.ops.functions.assert_valid(op)
+    # Skip differentiation for test cases that raise NaNs in gradients (known limitation of ``MottonenStatePreparation``).
+    qml.ops.functions.assert_valid(op, skip_differentiation=skip_diff)
 
 
 def test_repr():

tests/templates/test_subroutines/test_qubitization.py

@@ -65,12 +65,18 @@ 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", "skip_diff"),
+    [
+        (qml.dot([0.1, -0.3], [qml.X(2), qml.Z(3)]), [0], False),
+        (qml.dot([0.1, -0.3, -0.3], [qml.X(0), qml.Z(1), qml.Y(0) @ qml.Z(2)]), [3, 4], True),
+    ],
+)
+def test_standard_validity(lcu, control, skip_diff):
     """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])
-    qml.ops.functions.assert_valid(op)
+    op = qml.Qubitization(lcu, control)
+    # Skip differentiation for test cases that raise NaNs in gradients (known limitation of ``MottonenStatePreparation``).
+    qml.ops.functions.assert_valid(op, skip_differentiation=skip_diff)
 
 
 @pytest.mark.usefixtures("use_legacy_and_new_opmath")