Repository: PennyLaneAI/pennylane (python)

Issue:
[BUG] Adjoint and backpropagation differentiation gives different results for custom qml.Hermitian observables
### Expected behavior

A QML model with a custom `qml.Hermitian` observable should give the identical results when trained with either adjoint differentiation or backprop differentiation methods.

### Actual behavior

The two differentiation methods give different results: 
* The backprop trained algorithm learns and produces the expected accuracy.
* The adjoint trained algorithm does not learn at all, final predictions for a binary data set are always 0.5.

Again, as per the documentation, the two algorithms should produce identical results.

### Additional information

Here is the output produced by training with backprop and adjoint, respectively.

[backprop_output.txt](https://github.com/PennyLaneAI/pennylane/files/8463642/backprop_output.txt)
[adjoint_output.txt](https://github.com/PennyLaneAI/pennylane/files/8463643/adjoint_output.txt)

If we change to a default observable (line 35-36 in the code below) such as PauliZ we observe that adjoint and backprop give identical results, as you can see below.

[adjoint_output_pauliZ.txt](https://github.com/PennyLaneAI/pennylane/files/8463662/adjoint_output_pauliZ.txt)
[backprop_output_pauliZ.txt](https://github.com/PennyLaneAI/pennylane/files/8463663/backprop_output_pauliZ.txt)

This issue was first discussed here: https://discuss.pennylane.ai/t/adjoint-and-backpropagation-differentiation-methods-give-different-results/1731 .

### Source code

```shell
# Minimal example that reproduces the issue for the backprop and adjoint discrepancy.
import torch.nn as nn
import numpy as np
from sklearn.datasets import make_moons
import matplotlib.pyplot as plt
import torch
import pennylane as qml
import sys
from time import perf_counter

class Model(nn.Module):
    def __init__(self, dev, diff_method="backprop"):
        super().__init__()

        self.cnet_in = self.get_cnet_in()
        self.qcircuit = qml.qnode(dev, interface="torch", 
                                  diff_method=diff_method)(self.qnode)
        
        weight_shape = {"weights":(2,)}
        self.qlayer = qml.qnn.TorchLayer(self.qcircuit, weight_shape)

    def get_cnet_in(self):
        layers = [nn.Linear(2,16), nn.ReLU(True), nn.Linear(16,2), nn.Sigmoid()]
        return nn.Sequential(*layers)     
    
    def qnode(self, inputs, weights):
        # Data encoding:
        for x in range(len(inputs)):
            qml.Hadamard(x)
            qml.RZ(2.0 * inputs[x], wires=x)
        # Trainable part:
        qml.CNOT(wires=[0,1])
        qml.RY(weights[0], wires=0)
        qml.RY(weights[1], wires=1)
        o = [[1], [0]] * np.conj([[1], [0]]).T
        return qml.expval(qml.Hermitian(o, wires=[0]))
        #return qml.expval(qml.PauliY(wires=0))

    def forward(self, x):
        x = self.cnet_in(x)
        x = self.qlayer(x)
        return x

def train(X, y, dev_name, diff_method):
    
    dev = qml.device(dev_name, wires=2, shots=None)
    model  = Model(dev, diff_method)
    
    # Train the model
    opt = torch.optim.Adam(model.parameters(), lr=0.01)
    loss = torch.nn.MSELoss()

    X = torch.tensor(X, requires_grad=False).float()
    y = torch.Tensor(y).float()

    batch_size = 5
    batches = 200 // batch_size

    data_loader = torch.utils.data.DataLoader(
        list(zip(X, y)), batch_size=batch_size, shuffle=True, drop_last=True
    )

    epochs = 20
    avg_loss = []
    for epoch in range(epochs):

        running_loss = 0
        x = 0

        for xs, ys in data_loader:
            opt.zero_grad()

            loss_evaluated = loss(model(xs), ys)
            loss_evaluated.backward()
            opt.step()
            #print('Weights:', model.state_dict()['qlayer.weights'])
            running_loss += loss_evaluated
        loss_value = running_loss / batches
        avg_loss.append(loss_value.detach().numpy())
        print("Average loss over epoch {}: {:.10f}".format(epoch + 1, loss_value))

    y_pred = model(X)
    print(y_pred)
    predictions = (y_pred >= 0.5).type(torch.uint8)
    correct = [1 if p == p_true else 0 for p, p_true in zip(predictions, y)]
    accuracy = sum(correct) / len(correct)
    print(f"Accuracy: {accuracy * 100}%")
    
    return avg_loss


if __name__ == "__main__":
    torch.manual_seed(42)
    np.random.seed(42)
    X, y = make_moons(200)
    begin_time = perf_counter()
    avg_loss = train(X, y, str(sys.argv[1]), str(sys.argv[2]))
    end_time = perf_counter()
    runtime = end_time-begin_time
    print(f'Runtime: {runtime:.2e} s or {(runtime/60):.2e} min.')
```


### Tracebacks

_No response_

### System information

```shell
WARNING: pip is being invoked by an old script wrapper. This will fail in a future version of pip.
Please see https://github.com/pypa/pip/issues/5599 for advice on fixing the underlying issue.
To avoid this problem you can invoke Python with '-m pip' instead of running pip directly.
Name: PennyLane
Version: 0.20.0
Summary: PennyLane is a Python quantum machine learning library by Xanadu Inc.
Home-page: https://github.com/XanaduAI/pennylane
Author:
Author-email:
License: Apache License 2.0
Location: /home/vabelis/anaconda3/envs/ae_qml_pnl/lib/python3.8/site-packages
Requires: scipy, semantic-version, autograd, autoray, appdirs, pennylane-lightning, toml, numpy, cachetools, networkx
Required-by: PennyLane-qiskit, PennyLane-Lightning
Platform info:           Linux-3.10.0-1160.36.2.el7.x86_64-x86_64-with-glibc2.17
Python version:          3.8.12
Numpy version:           1.22.2
Scipy version:           1.7.3
Installed devices:
- default.gaussian (PennyLane-0.20.0)
- default.mixed (PennyLane-0.20.0)
- default.qubit (PennyLane-0.20.0)
- default.qubit.autograd (PennyLane-0.20.0)
- default.qubit.jax (PennyLane-0.20.0)
- default.qubit.tf (PennyLane-0.20.0)
- default.qubit.torch (PennyLane-0.20.0)
- qiskit.aer (PennyLane-qiskit-0.21.0)
- qiskit.basicaer (PennyLane-qiskit-0.21.0)
- qiskit.ibmq (PennyLane-qiskit-0.21.0)
- qiskit.ibmq.circuit_runner (PennyLane-qiskit-0.21.0)
- qiskit.ibmq.sampler (PennyLane-qiskit-0.21.0)
- lightning.qubit (PennyLane-Lightning-0.22.1)
```


### Existing GitHub issues

- [X] I have searched existing GitHub issues to make sure the issue does not already exist.

Pull request description:
Fix `diff_method="adjoint"` bug with `qml.Hermitian` (inc. adding `tape.get_parameters` `operations_only` arg)
**Context:**

`diff_method="adjoint"` is a device differentiation method and the main logic is implemented in the `QubitDevice.adjoint_jacobian` method. The implementation depends on a private `tape._par_info` dictionary and the `tape.trainable_params` list of trainable parameter indices.

The implementation doesn't, however, consider the case when measurements contain parametrized observables such as `qml.Hermitian` that are different than the disallowed `qml.Hamiltonian` observable.

**Description of the Change:**

* Add a new `circuit_only` argument to the `tape.get_parameters` method such that parameters affecting the circuit can be obtained and parameters to measured observables can be excluded;
* Changes the logic of `adjoint_jacobian` such that it filters out entries from `tape._par_info` that are possible coming from measurements (checked via `hasattr(tape._par_info[k]["op"], "return_type")`);
* Warns if parameters to measured observables were marked as trainable with the adjoint diff method;
* Updates the doc of `adjoint_jacobian` such that the page mentions that no parametrized observable is supported.

**Benefits:**
* #2434 is solved:
```
Average loss over epoch 20: 0.0661188588
<Output tensor as expected>
Accuracy: 97.5%
```
* The case of computing gradients of circuits with parametrized measured observables is cleared in the docs.

**Possible Drawbacks:**
N/A

**Related GitHub Issues:**
Closes #2434.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 588621e9d..e233745f1 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -12,6 +12,9 @@
 
 <h3>Improvements</h3>
 
+* Introduced an `operations_only` argument to the `tape.get_parameters` method.
+  [(#2543)](https://github.com/PennyLaneAI/pennylane/pull/2543)
+
 * The `gradients` module now uses faster subroutines and uniform
   formats of gradient rules.
   [(#2452)](https://github.com/XanaduAI/pennylane/pull/2452)
@@ -71,6 +74,10 @@
 
 <h3>Bug fixes</h3>
 
+* Fixed a bug for `diff_method="adjoint"` where incorrect gradients were
+  computed for QNodes with parametrized observables (e.g., `qml.Hermitian`).
+  [(#2543)](https://github.com/PennyLaneAI/pennylane/pull/2543)
+
 * Fixed a bug where `QNGOptimizer` did not work with operators
   whose generator was a Hamiltonian.
   [(#2524)](https://github.com/PennyLaneAI/pennylane/pull/2524)
diff --git a/pennylane/_qubit_device.py b/pennylane/_qubit_device.py
index ca9be8b12..2bfa5a802 100644
--- a/pennylane/_qubit_device.py
+++ b/pennylane/_qubit_device.py
@@ -939,7 +939,8 @@ class QubitDevice(Device):
 
             * Only expectation values are supported as measurements.
 
-            * Does not work for Hamiltonian observables.
+            * Does not work for parametrized observables like
+              :class:`~.Hamiltonian` or :class:`~.Hermitian`.
 
         Args:
             tape (.QuantumTape): circuit that the function takes the gradient of
@@ -1014,13 +1015,29 @@ class QubitDevice(Device):
                 if op.name not in ("QubitStateVector", "BasisState", "Snapshot"):
                     expanded_ops.append(op)
 
-        jac = np.zeros((len(tape.observables), len(tape.trainable_params)))
+        trainable_params = []
+        for k in tape.trainable_params:
+            # pylint: disable=protected-access
+            if hasattr(tape._par_info[k]["op"], "return_type"):
+                warnings.warn(
+                    "Differentiating with respect to the input parameters of "
+                    f"{tape._par_info[k]['op'].name} is not supported with the "
+                    "adjoint differentiation method. Gradients are computed "
+                    "only with regards to the trainable parameters of the circuit.\n\n Mark "
+                    "the parameters of the measured observables as non-trainable "
+                    "to silence this warning.",
+                    UserWarning,
+                )
+            else:
+                trainable_params.append(k)
+
+        jac = np.zeros((len(tape.observables), len(trainable_params)))
 
-        param_number = len(tape._par_info) - 1  # pylint: disable=protected-access
-        trainable_param_number = len(tape.trainable_params) - 1
+        param_number = len(tape.get_parameters(trainable_only=False, operations_only=True)) - 1
+        trainable_param_number = len(trainable_params) - 1
         for op in expanded_ops:
 
-            if (op.grad_method is not None) and (param_number in tape.trainable_params):
+            if (op.grad_method is not None) and (param_number in trainable_params):
                 d_op_matrix = operation_derivative(op)
 
             op.inv()
@@ -1029,7 +1046,7 @@ class QubitDevice(Device):
             ket = self._apply_operation(ket, op)
 
             if op.grad_method is not None:
-                if param_number in tape.trainable_params:
+                if param_number in trainable_params:
 
                     ket_temp = self._apply_unitary(ket, d_op_matrix, op.wires)
 
diff --git a/pennylane/tape/tape.py b/pennylane/tape/tape.py
index 30fa6969a..20b2f3110 100644
--- a/pennylane/tape/tape.py
+++ b/pennylane/tape/tape.py
@@ -795,7 +795,9 @@ class QuantumTape(AnnotatedQueue):
         p_idx = info["p_idx"]
         return op, p_idx
 
-    def get_parameters(self, trainable_only=True, **kwargs):  # pylint:disable=unused-argument
+    def get_parameters(
+        self, trainable_only=True, operations_only=False, **kwargs
+    ):  # pylint:disable=unused-argument
         """Return the parameters incident on the tape operations.
 
         The returned parameters are provided in order of appearance
@@ -803,6 +805,8 @@ class QuantumTape(AnnotatedQueue):
 
         Args:
             trainable_only (bool): if True, returns only trainable parameters
+            operations_only (bool): if True, returns only the parameters of the
+                operations excluding parameters to observables of measurements
 
         **Example**
 
@@ -838,6 +842,9 @@ class QuantumTape(AnnotatedQueue):
 
         for p_idx in iterator:
             op = self._par_info[p_idx]["op"]
+            if operations_only and hasattr(op, "return_type"):
+                continue
+
             op_idx = self._par_info[p_idx]["p_idx"]
             params.append(op.data[op_idx])
         return params


Test patch:
diff --git a/tests/interfaces/test_jax_qnode.py b/tests/interfaces/test_jax_qnode.py
index 51adda992..f46440514 100644
--- a/tests/interfaces/test_jax_qnode.py
+++ b/tests/interfaces/test_jax_qnode.py
@@ -1336,6 +1336,9 @@ class TestJIT:
         if diff_method == "backprop":
             pytest.skip("Backpropagation is unsupported if shots > 0.")
 
+        if diff_method == "adjoint" and mode == "forward":
+            pytest.skip("Computing the gradient for Hermitian is not supported with adjoint.")
+
         projector = np.array(qml.matrix(qml.PauliZ(0) @ qml.PauliZ(1)))
 
         @jax.jit
diff --git a/tests/tape/test_tape.py b/tests/tape/test_tape.py
index e514b32dc..d309c71e2 100644
--- a/tests/tape/test_tape.py
+++ b/tests/tape/test_tape.py
@@ -606,6 +606,20 @@ class TestParameters:
 
         return tape, params
 
+    @pytest.fixture
+    def make_tape_with_hermitian(self):
+        params = [0.432, 0.123, 0.546, 0.32, 0.76]
+        hermitian = qml.numpy.eye(2, requires_grad=False)
+
+        with QuantumTape() as tape:
+            qml.RX(params[0], wires=0)
+            qml.Rot(*params[1:4], wires=0)
+            qml.CNOT(wires=[0, "a"])
+            qml.RX(params[4], wires=4)
+            qml.expval(qml.Hermitian(hermitian, wires="a"))
+
+        return tape, params, hermitian
+
     def test_parameter_processing(self, make_tape):
         """Test that parameters are correctly counted and processed"""
         tape, params = make_tape
@@ -613,6 +627,20 @@ class TestParameters:
         assert tape.trainable_params == list(range(len(params)))
         assert tape.get_parameters() == params
 
+    @pytest.mark.parametrize("operations_only", [False, True])
+    def test_parameter_processing_operations_only(self, make_tape_with_hermitian, operations_only):
+        """Test the operations_only flag for getting the parameters on a tape with
+        qml.Hermitian is measured"""
+        tape, circuit_params, hermitian = make_tape_with_hermitian
+        num_all_params = len(circuit_params) + 1  # + 1 for hermitian
+        assert tape.num_params == num_all_params
+        assert tape.trainable_params == list(range(num_all_params))
+        assert (
+            tape.get_parameters(operations_only=operations_only) == circuit_params
+            if operations_only
+            else circuit_params + [hermitian]
+        )
+
     def test_set_trainable_params(self, make_tape):
         """Test that setting trainable parameters works as expected"""
         tape, params = make_tape
diff --git a/tests/test_qubit_device_adjoint_jacobian.py b/tests/test_qubit_device_adjoint_jacobian.py
index 6cfd600b2..afbb58eae 100644
--- a/tests/test_qubit_device_adjoint_jacobian.py
+++ b/tests/test_qubit_device_adjoint_jacobian.py
@@ -250,6 +250,52 @@ class TestAdjointJacobian:
 
         assert np.allclose(dM1, dM2, atol=tol, rtol=0)
 
+    def test_gradient_of_tape_with_hermitian(self, tol):
+        """Test that computing the gradient of a tape that obtains the
+        expectation value of a Hermitian operator works correctly."""
+        dev = qml.device("default.qubit", wires=3)
+
+        a, b, c = [0.5, 0.3, -0.7]
+
+        def ansatz(a, b, c):
+            qml.RX(a, wires=0)
+            qml.RX(b, wires=1)
+            qml.RX(c, wires=2)
+            qml.CNOT(wires=[0, 1])
+            qml.CNOT(wires=[1, 2])
+
+        mx = qml.matrix(qml.PauliX(0) @ qml.PauliY(2))
+        with qml.tape.QuantumTape() as tape:
+            ansatz(a, b, c)
+            qml.RX(a, wires=0)
+            qml.expval(qml.Hermitian(mx, wires=[0, 2]))
+
+        tape.trainable_params = {0, 1, 2}
+        res = dev.adjoint_jacobian(tape)
+
+        expected = [
+            np.cos(a) * np.sin(b) * np.sin(c),
+            np.cos(b) * np.sin(a) * np.sin(c),
+            np.cos(c) * np.sin(b) * np.sin(a),
+        ]
+        assert np.allclose(res, expected, atol=tol, rtol=0)
+
+    def test_trainable_hermitian_warns(self, tol):
+        """Test attempting to compute the gradient of a tape that obtains the
+        expectation value of a Hermitian operator emits a warning if the
+        parameters to Hermitian are trainable."""
+        dev = qml.device("default.qubit", wires=3)
+
+        mx = qml.matrix(qml.PauliX(0) @ qml.PauliY(2))
+        with qml.tape.QuantumTape() as tape:
+            qml.expval(qml.Hermitian(mx, wires=[0, 2]))
+
+        tape.trainable_params = {0}
+        with pytest.warns(
+            UserWarning, match="Differentiating with respect to the input parameters of Hermitian"
+        ):
+            res = dev.adjoint_jacobian(tape)
+
 
 class TestAdjointJacobianQNode:
     """Test QNode integration with the adjoint_jacobian method"""
@@ -478,3 +524,36 @@ class TestAdjointJacobianQNode:
         grad_backprop = jax.grad(qnode_backprop)(params1, params2)
 
         assert np.allclose(grad_adjoint, grad_backprop)
+
+    def test_gradient_of_qnode_with_hermitian(self, tol):
+        """Test that computing the gradient of a QNode that obtains the
+        expectation value of a Hermitian operator works correctly."""
+        dev = qml.device("default.qubit", wires=3)
+
+        def ansatz(a, b, c):
+            qml.RX(a, wires=0)
+            qml.RX(b, wires=1)
+            qml.RX(c, wires=2)
+            qml.CNOT(wires=[0, 1])
+            qml.CNOT(wires=[1, 2])
+
+        mx = qml.matrix(qml.PauliX(0) @ qml.PauliY(2))
+
+        @qml.qnode(dev, diff_method="adjoint")
+        def circ(a, b, c):
+            ansatz(a, b, c)
+            qml.RX(a, wires=0)
+            return qml.expval(qml.Hermitian(mx, wires=[0, 2]))
+
+        a = np.array(0.5, requires_grad=True)
+        b = np.array(0.3, requires_grad=True)
+        c = np.array(-0.7, requires_grad=True)
+
+        res = qml.grad(circ)(a, b, c)
+
+        expected = [
+            np.cos(a) * np.sin(b) * np.sin(c),
+            np.cos(b) * np.sin(a) * np.sin(c),
+            np.cos(c) * np.sin(b) * np.sin(a),
+        ]
+        assert np.allclose(res, expected, atol=tol, rtol=0)