Repository: PennyLaneAI/pennylane
Pull request #4820 (2023-11-10 14:39:08)

Fix parameter counting in `DefaultQubitLegacy` adjoint method
**Context:**
The adjoint method of `DefaultQubitLegacy` uses two incompatible ways to quantify tape parameters: It sets up the number of tape parameters correctly using `tape.get_parameters(trainable_only=False)` but checks whether an operation comes with a parameter via `grad_method is not None`. This lead to a bug where, for example, `QubitUnitary` disrupts the parameter counting.


**Description of the Change:**
As all operations with more than one parameter are being decomposed/raised an error upon anyways,  this PR changes the way to detect parametrized operations to `op.num_params==1`.

**Benefits:**
Bug fix: Correct adjoint method derivatives even when using `QubitUnitary` and other operations that have `grad_method=None` but have parameters.

**Possible Drawbacks:**
N/A

**Related GitHub Issues:**
N/A


dwierichs: @vincentmr This is the PR (sorry, took a moment to figure this out after all)
dwierichs: >Even though its currently not causing issues, mind fixing this as well in the new system:

@albi3ro Done :) 

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index a85ba4548..ac75e056a 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -75,6 +75,10 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* Fixes a bug where the adjoint method differentiation would fail if
+  an operation with `grad_method=None` that has a parameter is present.
+  [(#4820)](https://github.com/PennyLaneAI/pennylane/pull/4820)
+  
 * `MottonenStatePreparation` now raises an error if decomposing a broadcasted state vector.
   [(#4767)](https://github.com/PennyLaneAI/pennylane/pull/4767)
 
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index 7d4d7778d..a74bd8ae7 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -1669,7 +1669,7 @@ class QubitDevice(Device):
             adj_op = qml.adjoint(op)
             ket = self._apply_operation(ket, adj_op)
 
-            if op.grad_method is not None:
+            if op.num_params == 1:
                 if param_number in trainable_params:
                     d_op_matrix = operation_derivative(op)
                     ket_temp = self._apply_unitary(ket, d_op_matrix, op.wires)
diff --git a/pennylane/devices/qubit/adjoint_jacobian.py b/pennylane/devices/qubit/adjoint_jacobian.py
index b4251b6a7..79a712560 100644
--- a/pennylane/devices/qubit/adjoint_jacobian.py
+++ b/pennylane/devices/qubit/adjoint_jacobian.py
@@ -80,7 +80,7 @@ def adjoint_jacobian(tape: QuantumTape, state=None):
         adj_op = qml.adjoint(op)
         ket = apply_operation(adj_op, ket)
 
-        if op.grad_method is not None:
+        if op.num_params == 1:
             if param_number in tape.trainable_params:
                 d_op_matrix = operation_derivative(op)
                 ket_temp = apply_operation(qml.QubitUnitary(d_op_matrix, wires=op.wires), ket)
@@ -157,7 +157,7 @@ def adjoint_jvp(tape: QuantumTape, tangents: Tuple[Number], state=None):
         adj_op = qml.adjoint(op)
         ket = apply_operation(adj_op, ket)
 
-        if op.grad_method is not None:
+        if op.num_params == 1:
             if param_number in tape.trainable_params:
                 # don't do anything if the tangent is 0
                 if not np.allclose(tangents[trainable_param_number], 0):
@@ -232,7 +232,7 @@ def adjoint_vjp(tape: QuantumTape, cotangents: Tuple[Number], state=None):
         adj_op = qml.adjoint(op)
         ket = apply_operation(adj_op, ket)
 
-        if op.grad_method is not None:
+        if op.num_params == 1:
             if param_number in tape.trainable_params:
                 d_op_matrix = operation_derivative(op)
                 ket_temp = apply_operation(qml.QubitUnitary(d_op_matrix, wires=op.wires), ket)


Test patch:
diff --git a/tests/devices/qubit/test_adjoint_jacobian.py b/tests/devices/qubit/test_adjoint_jacobian.py
index 26e40f6ee..7600309c3 100644
--- a/tests/devices/qubit/test_adjoint_jacobian.py
+++ b/tests/devices/qubit/test_adjoint_jacobian.py
@@ -294,6 +294,23 @@ class TestAdjointJacobian:
         ]
         assert np.allclose(res, expected, atol=tol, rtol=0)
 
+    def test_with_nontrainable_parametrized(self):
+        """Test that a parametrized `QubitUnitary` is accounted for correctly
+        when it is not trainable."""
+
+        par = np.array(0.6)
+
+        ops = [
+            qml.RY(par, wires=0),
+            qml.QubitUnitary(np.eye(2), wires=0),
+        ]
+        qs = QuantumScript(ops, [qml.expval(qml.PauliZ(0))])
+        qs.trainable_params = [0]
+
+        grad_adjoint = adjoint_jacobian(qs)
+        expected = [-np.sin(par)]
+        assert np.allclose(grad_adjoint, expected)
+
 
 class TestAdjointJVP:
     """Test for adjoint_jvp"""
@@ -388,6 +405,24 @@ class TestAdjointJVP:
         expected = jac @ np.array(tangents)
         assert np.allclose(actual, expected, atol=tol)
 
+    def test_with_nontrainable_parametrized(self):
+        """Test that a parametrized `QubitUnitary` is accounted for correctly
+        when it is not trainable."""
+
+        par = np.array(0.6)
+        tangents = (0.45,)
+
+        ops = [
+            qml.RY(par, wires=0),
+            qml.QubitUnitary(np.eye(2), wires=0),
+        ]
+        qs = QuantumScript(ops, [qml.expval(qml.PauliZ(0))])
+        qs.trainable_params = [0]
+
+        jvp_adjoint = adjoint_jvp(qs, tangents)
+        expected = [-np.sin(par) * tangents[0]]
+        assert np.allclose(jvp_adjoint, expected)
+
 
 class TestAdjointVJP:
     """Test for adjoint_vjp"""
@@ -484,3 +519,21 @@ class TestAdjointVJP:
         jac = np.array([[-np.sin(x), 0], [0, -np.cos(y)], [np.cos(x), 0]])
         expected = np.array(cotangents) @ jac
         assert np.allclose(actual, expected, atol=tol)
+
+    def test_with_nontrainable_parametrized(self):
+        """Test that a parametrized `QubitUnitary` is accounted for correctly
+        when it is not trainable."""
+
+        par = np.array(0.6)
+        cotangents = (0.45,)
+
+        ops = [
+            qml.RY(par, wires=0),
+            qml.QubitUnitary(np.eye(2), wires=0),
+        ]
+        qs = QuantumScript(ops, [qml.expval(qml.PauliZ(0))])
+        qs.trainable_params = [0]
+
+        vjp_adjoint = adjoint_vjp(qs, cotangents)
+        expected = [-np.sin(par) * cotangents[0]]
+        assert np.allclose(vjp_adjoint, expected)
diff --git a/tests/gradients/core/test_adjoint_diff.py b/tests/gradients/core/test_adjoint_diff.py
index 7e85b1405..65e726303 100644
--- a/tests/gradients/core/test_adjoint_diff.py
+++ b/tests/gradients/core/test_adjoint_diff.py
@@ -419,3 +419,21 @@ class TestAdjointJacobian:
             assert all(isinstance(g, np.ndarray) for g in expected)
 
             assert np.allclose(grad_D[i], expected)
+
+    def test_with_nontrainable_parametrized(self):
+        """Test that a parametrized `QubitUnitary` is accounted for correctly
+        when it is not trainable."""
+
+        dev = qml.device("default.qubit.legacy", wires=1)
+        par = np.array(0.6)
+
+        def circuit(x):
+            qml.RY(x, wires=0)
+            qml.QubitUnitary(np.eye(2, requires_grad=False), wires=0)
+            return qml.expval(qml.PauliZ(0))
+
+        circ = qml.QNode(circuit, dev, diff_method="adjoint")
+        grad_adjoint = qml.jacobian(circ)(par)
+        circ = qml.QNode(circuit, dev, diff_method="parameter-shift")
+        grad_psr = qml.jacobian(circ)(par)
+        assert np.allclose(grad_adjoint, grad_psr)