Repository: PennyLaneAI/pennylane (python)

Issue:
[BUG] Differentiating `PrepSelPrep` using parameter-shift returns 0s
### Expected behavior

`PrepSelPrep` can be differentiated properly

### Actual behavior

```python
import pennylane as qml

@qml.qnode(qml.device("default.qubit"), diff-method="parameter-shift")
def circuit(params):
    H = qml.Hamiltonian(params, [qml.Z(2), qml.X(1) @ qml.X(2)])
    qml.PrepSelPrep(H, control=[0])
    return qml.expval(qml.PauliZ(0))

>>> x = qml.numpy.array([0.25, 0.75])
>>> qml.jacobian(circuit)(x)
array([0.,  0.])
```

### Additional information

The issue was discovered as part of https://github.com/PennyLaneAI/pennylane/pull/6282, and the relevant tests in `test_prepselprep.py` and `test_quibitization.py` are xfailed, when fixed, remove the xfail from the relevant tests.

### Source code

_No response_

### Tracebacks

_No response_

### System information

```shell
Name: PennyLane
Version: 0.39.0.dev10
Summary: PennyLane is a cross-platform Python library for quantum computing, quantum machine learning, and quantum chemistry. Train a quantum computer the same way as a neural network.
Home-page: https://github.com/PennyLaneAI/pennylane
Author:
Author-email:
License: Apache License 2.0
Location: /Users/astral.cai/Workspace/pennylane/venv/lib/python3.10/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, packaging, pennylane-lightning, requests, rustworkx, scipy, toml, typing-extensions
Required-by: PennyLane-Catalyst, PennyLane_Lightning, PennyLane_Lightning_Kokkos

Platform info:           macOS-15.0-arm64-arm-64bit
Python version:          3.10.14
Numpy version:           1.26.4
Scipy version:           1.12.0
Installed devices:
- default.clifford (PennyLane-0.39.0.dev19)
- default.gaussian (PennyLane-0.39.0.dev19)
- default.mixed (PennyLane-0.39.0.dev19)
- default.qubit (PennyLane-0.39.0.dev19)
- default.qutrit (PennyLane-0.39.0.dev19)
- default.qutrit.mixed (PennyLane-0.39.0.dev19)
- default.tensor (PennyLane-0.39.0.dev19)
- null.qubit (PennyLane-0.39.0.dev19)
- reference.qubit (PennyLane-0.39.0.dev19)
- lightning.qubit (PennyLane_Lightning-0.39.0.dev27)
- lightning.kokkos (PennyLane_Lightning_Kokkos-0.38.0)
- nvidia.custatevec (PennyLane-Catalyst-0.9.0.dev15)
- nvidia.cutensornet (PennyLane-Catalyst-0.9.0.dev15)
- oqc.cloud (PennyLane-Catalyst-0.9.0.dev15)
- softwareq.qpp (PennyLane-Catalyst-0.9.0.dev15)
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.

Pull request:
Fix `PrepSelPrep` differentiability with `parameter-shift`
**Context:**

Prior to this fix, differentiating `PrepSelPrep` with `diff_method="parameter-shift"` resulted in zeros (which was inconsistent with `backprop`),

```python
@qml.qnode(qml.device("default.qubit"), diff-method="parameter-shift")
def circuit(params):
    H = qml.Hamiltonian(params, [qml.Z(2), qml.X(1) @ qml.X(2)])
    qml.PrepSelPrep(H, control=[0])
    return qml.expval(qml.PauliZ(0))

>>> x = qml.numpy.array([0.25, 0.75])
>>> qml.jacobian(circuit)(x)
array([0.,  0.])
```

**Description of the Change:**

One line change to add `PrepSelPrep.grad_method=None` 😺. This results in,

```python
@qml.qnode(qml.device("default.qubit"), diff-method="parameter-shift")
def circuit(params):
    H = qml.Hamiltonian(params, [qml.Z(2), qml.X(1) @ qml.X(2)])
    qml.PrepSelPrep(H, control=[0])
    return qml.expval(qml.PauliZ(0))

>>> x = qml.numpy.array([0.25, 0.75])
>>> qml.jacobian(circuit)(x)
array([-1.5, 0.5])
```

_Reason this works:_ It seems that parameter-shift tries to update the parameters and fails to do so. Setting `grad_method=None` forces it to decompose the operator before differentiating it's data. This does suggest some areas of improvement for `Operator.grad_method`.

The `xfails` were also removed from the relevant tests. Also, certain test cases would result in `NaN`s upon differentiation with *any* method, so I skipped those cases as they are a *known* limitation of `MottonenStatePreparation`.

**Benefits:** Gradient results are now consistent with `backprop`.

**Possible Drawbacks:** None.

**Related GitHub Issues:** Fixes #6331 

[sc-74920]


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 08edfa3f5..0c3bff398 100644
--- a/doc/releases/changelog-dev.md
+++ b/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)
 
diff --git a/pennylane/templates/subroutines/prepselprep.py b/pennylane/templates/subroutines/prepselprep.py
index d9aa0b11a..6d8a7b660 100644
--- a/pennylane/templates/subroutines/prepselprep.py
+++ b/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()


Test patch:
diff --git a/tests/templates/test_subroutines/test_prepselprep.py b/tests/templates/test_subroutines/test_prepselprep.py
index b6fdbf00b..cf0849f98 100644
--- a/tests/templates/test_subroutines/test_prepselprep.py
+++ b/tests/templates/test_subroutines/test_prepselprep.py
@@ -23,27 +23,28 @@ import pytest
 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():
diff --git a/tests/templates/test_subroutines/test_qubitization.py b/tests/templates/test_subroutines/test_qubitization.py
index 7649a0a7d..66e40066e 100644
--- a/tests/templates/test_subroutines/test_qubitization.py
+++ b/tests/templates/test_subroutines/test_qubitization.py
@@ -65,12 +65,18 @@ def test_operator_definition_qpe(hamiltonian):
     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")