Repository: PennyLaneAI/pennylane (python)

Issue:
[BUG] Issue with the gradients of the expectation value of an operator
### Expected behavior

Hi, I'm trying to compute the gradient of the circuit with respect to the parameters on the operator $\theta_{i,i+1}(\sigma^+_i\sigma^-_{i+1} + \sigma^+_{i+1}\sigma^-_i)$ where $\theta$ is my variable. When I compute the expectation value of the circuit, I can see that individual terms of my operator are differentiable, e.g. for $\theta_0(\sigma^+_0\sigma^-_1 + \sigma^+_1\sigma^-_0)$ I can get a gradient with respect to $\theta_0$ but as soon as I add other terms I get None which I was expecting to get the gradient for every term.

### Actual behavior

Gradient returns ``None``.

### Additional information

I get zero for the gradient when I use autograd or Jax.

### Source code

```shell
import pennylane as qml
import tensorflow as tf
import pennylane.numpy as np

# here s_prod or sum is not necessary +, @ and * works the same
sig_plus = lambda w: qml.s_prod(0.5, (qml.sum(qml.PauliX(w) , qml.s_prod(1j, qml.PauliY(w)))))
sig_minus = lambda w: qml.s_prod(0.5, (qml.sum(qml.PauliX(w) , qml.s_prod(-1j, qml.PauliY(w)))))
term = lambda w: qml.sum(qml.prod(sig_plus(w) , sig_minus(w + 1)) , qml.prod(sig_plus(w + 1), sig_minus(w)))

@qml.qnode(qml.device("default.qubit.tf", wires=4), interface="tf")
def circ(phi, theta):
    # This portion of the circuit is random not important
    qml.RY(phi[0], 0)
    qml.RY(phi[1], 1)
    qml.RY(phi[2], 2)
    qml.CNOT([0,1])
    qml.CNOT([1,2])
    ###########################
    return qml.expval(theta[0]*term(0) + theta[1]*term(1) ) # only 0 or only 1 works just fine, also tried with qml.sum and qml.s_prod

theta = tf.Variable(np.random.uniform(0,1,3))
phi = tf.Variable(np.random.uniform(0,1,3))

def cost(phi, theta):
    return circ(phi, theta)

with tf.GradientTape() as tape:
    loss = cost(phi, theta)
tape.gradient(loss, [phi, theta])

#[<tf.Tensor: shape=(3,), dtype=float64, numpy=array([0.00131378, 0.00458278, 0.00750094])>,
# None]
```


### Tracebacks

```shell
No error has been generated
```


### System information

```shell
TF version 2.11.0
TF compiler version Apple LLVM 13.1.6 (clang-1316.0.21.1)

Name: PennyLane
Version: 0.30.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: /Users/jackaraz/packages/miniconda3/lib/python3.9/site-packages
Requires: appdirs, autograd, autoray, cachetools, networkx, numpy, pennylane-lightning, requests, rustworkx, scipy, semantic-version, toml
Required-by: field-theory, PennyLane-Lightning, PennyLane-qiskit, qhbm

Platform info:           macOS-13.3.1-arm64-arm-64bit
Python version:          3.9.16
Numpy version:           1.23.5
Scipy version:           1.10.0
Installed devices:
- default.gaussian (PennyLane-0.30.0)
- default.mixed (PennyLane-0.30.0)
- default.qubit (PennyLane-0.30.0)
- default.qubit.autograd (PennyLane-0.30.0)
- default.qubit.jax (PennyLane-0.30.0)
- default.qubit.tf (PennyLane-0.30.0)
- default.qubit.torch (PennyLane-0.30.0)
- default.qutrit (PennyLane-0.30.0)
- null.qubit (PennyLane-0.30.0)
- lightning.qubit (PennyLane-Lightning-0.30.0)
- qiskit.aer (PennyLane-qiskit-0.28.0)
- qiskit.basicaer (PennyLane-qiskit-0.28.0)
- qiskit.ibmq (PennyLane-qiskit-0.28.0)
- qiskit.ibmq.circuit_runner (PennyLane-qiskit-0.28.0)
- qiskit.ibmq.sampler (PennyLane-qiskit-0.28.0)
```


### Existing GitHub issues

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

Pull request description:
support trainable Sum observables
**Context:**
Sum observables should be able to do anything that Hamiltonians can do. It turns out that computing the gradient of parameters on Sum observables did not work.

**Description of the Change:**
- Updates DefaultQubit to compute the expectation of Sums using the new `qml.devices.qubit.measure` to maintain trainability
- Updates `qml.devices.qubit.measure.get_measurement_function` to _always_ return `sum_of_terms_method` for Sums in backprop mode
 
A few little changes had to be made along the way for this to work with all autodiff frameworks:
- `obs.eigvals()` would be returned as a tracer if the parameters are trainable, but we called `np.array()` on it. Now I use `qml.math.asarray()` instead to avoid breaking trainability
- Torch tracers can't do scalar multiplication with numpy arrays, so that broke `SProd.eigvals` with a traced scalar. I use `qml.math.convert_like` to fix this, because it has custom code for exactly this

**Benefits:**
- We can compute the gradient of Sum observables!
- Very little code had to be added

**Possible Drawbacks:**
- only backprop support for now
- finite-shots does not work

**Related GitHub Issues:**
Fixes #4199 

**Future work**
- Support the things not yet supported (see the drawbacks section)
- Improve the choice made in get_measurement_function. Eg. if the Sum has no trainable parameters on it, this always-sum-of-terms solution might not be necessary

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index b4de43a35..07084e2f4 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -438,6 +438,9 @@
   An if conditional was intended to prevent divide by zero errors but the division was by the sine of the argument so any multiple of $\pi$ should trigger the conditional, but it was only checking if the argument was 0. Example: `qml.Rot(2.3, 2.3, 2.3)`
   [(#4210)](https://github.com/PennyLaneAI/pennylane/pull/4210)
 
+* Allow for `Sum` observables with trainable parameters.
+  [(#4251)](https://github.com/PennyLaneAI/pennylane/pull/4251)
+
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):
diff --git a/pennylane/devices/default_qubit.py b/pennylane/devices/default_qubit.py
index f39c5c183..35a0a5da1 100644
--- a/pennylane/devices/default_qubit.py
+++ b/pennylane/devices/default_qubit.py
@@ -28,9 +28,12 @@ from scipy.sparse import csr_matrix
 
 import pennylane as qml
 from pennylane import BasisState, DeviceError, QubitDevice, QubitStateVector, Snapshot
+from pennylane.devices.qubit import measure
 from pennylane.operation import Operation
+from pennylane.ops import Sum
 from pennylane.ops.qubit.attributes import diagonal_in_z_basis
 from pennylane.pulse import ParametrizedEvolution
+from pennylane.measurements import ExpectationMP
 from pennylane.typing import TensorLike
 from pennylane.wires import WireError
 
@@ -567,6 +570,12 @@ class DefaultQubit(QubitDevice):
             Hamiltonian is not NumPy or Autograd
 
         """
+        # intercept Sums
+        if isinstance(observable, Sum) and not self.shots:
+            return measure(
+                ExpectationMP(observable.map_wires(self.wire_map)), self._pre_rotated_state
+            )
+
         # intercept other Hamiltonians
         # TODO: Ideally, this logic should not live in the Device, but be moved
         # to a component that can be re-used by devices as needed.
diff --git a/pennylane/devices/qubit/measure.py b/pennylane/devices/qubit/measure.py
index 754c1b158..e30a0c809 100644
--- a/pennylane/devices/qubit/measure.py
+++ b/pennylane/devices/qubit/measure.py
@@ -135,18 +135,24 @@ def get_measurement_function(
             if measurementprocess.obs.name == "SparseHamiltonian":
                 return csr_dot_products
 
-            if isinstance(measurementprocess.obs, Hamiltonian) or (
-                isinstance(measurementprocess.obs, Sum)
-                and measurementprocess.obs.has_overlapping_wires
-                and len(measurementprocess.obs.wires) > 7
-            ):
-                # Use tensor contraction for `Sum` expectation values with non-commuting summands
-                # and 8 or more wires as it's faster than using eigenvalues.
-
+            backprop_mode = math.get_interface(state) != "numpy"
+            if isinstance(measurementprocess.obs, Hamiltonian):
                 # need to work out thresholds for when its faster to use "backprop mode" measurements
-                backprop_mode = math.get_interface(state) != "numpy"
                 return sum_of_terms_method if backprop_mode else csr_dot_products
 
+            if isinstance(measurementprocess.obs, Sum):
+                if backprop_mode:
+                    # always use sum_of_terms_method for Sum observables in backprop mode
+                    return sum_of_terms_method
+                if (
+                    measurementprocess.obs.has_overlapping_wires
+                    and len(measurementprocess.obs.wires) > 7
+                ):
+                    # Use tensor contraction for `Sum` expectation values with non-commuting summands
+                    # and 8 or more wires as it's faster than using eigenvalues.
+
+                    return csr_dot_products
+
         if measurementprocess.obs is None or measurementprocess.obs.has_diagonalizing_gates:
             return state_diagonalizing_gates
 
diff --git a/pennylane/measurements/measurements.py b/pennylane/measurements/measurements.py
index 92acc4d6c..5391b8605 100644
--- a/pennylane/measurements/measurements.py
+++ b/pennylane/measurements/measurements.py
@@ -23,8 +23,6 @@ from abc import ABC, abstractmethod
 from enum import Enum
 from typing import Sequence, Tuple, Optional
 
-import numpy as np
-
 import pennylane as qml
 from pennylane.operation import Operator
 from pennylane.wires import Wires
@@ -154,7 +152,7 @@ class MeasurementProcess(ABC):
             if obs is not None:
                 raise ValueError("Cannot set the eigenvalues if an observable is provided.")
 
-            self._eigvals = np.array(eigvals)
+            self._eigvals = qml.math.asarray(eigvals)
 
         # TODO: remove the following lines once devices
         # have been refactored to accept and understand receiving
diff --git a/pennylane/ops/op_math/sprod.py b/pennylane/ops/op_math/sprod.py
index 80fcd0812..a59f4a887 100644
--- a/pennylane/ops/op_math/sprod.py
+++ b/pennylane/ops/op_math/sprod.py
@@ -229,7 +229,10 @@ class SProd(ScalarSymbolicOp):
         Returns:
             array: array containing the eigenvalues of the operator.
         """
-        return self.scalar * self.base.eigvals()
+        base_eigs = self.base.eigvals()
+        if qml.math.get_interface(self.scalar) == "torch" and self.scalar.requires_grad:
+            base_eigs = qml.math.convert_like(base_eigs, self.scalar)
+        return self.scalar * base_eigs
 
     def sparse_matrix(self, wire_order=None):
         """Computes, by default, a `scipy.sparse.csr_matrix` representation of this Tensor.


Test patch:
diff --git a/tests/devices/test_default_qubit.py b/tests/devices/test_default_qubit.py
index 405a29233..85a2ef867 100644
--- a/tests/devices/test_default_qubit.py
+++ b/tests/devices/test_default_qubit.py
@@ -2361,6 +2361,75 @@ class TestHamiltonianSupport:
         assert qml.equal(call_args.measurements[0], qml.expval(qml.PauliX(0)))
 
 
+class TestSumSupport:
+    """Tests for custom Sum support in DefaultQubit."""
+
+    expected_grad = [-np.sin(1.3), np.cos(1.3)]
+
+    @staticmethod
+    def circuit(y, z):
+        qml.RX(1.3, 0)
+        return qml.expval(
+            qml.sum(
+                qml.s_prod(y, qml.PauliY(0)),
+                qml.s_prod(z, qml.PauliZ(0)),
+            )
+        )
+
+    def test_super_expval_not_called(self, mocker):
+        """Tests basic expval result, and ensures QubitDevice.expval is not called."""
+        dev = qml.device("default.qubit", wires=1)
+        spy = mocker.spy(qml.QubitDevice, "expval")
+        obs = qml.sum(qml.s_prod(0.1, qml.PauliX(0)), qml.s_prod(0.2, qml.PauliZ(0)))
+        assert np.isclose(dev.expval(obs), 0.2)
+        spy.assert_not_called()
+
+    @pytest.mark.autograd
+    def test_trainable_autograd(self):
+        """Tests that coeffs passed to a sum are trainable with autograd."""
+        dev = qml.device("default.qubit", wires=1)
+        qnode = qml.QNode(self.circuit, dev, interface="autograd")
+        y, z = np.array([1.1, 2.2])
+        actual = qml.grad(qnode)(y, z)
+        assert np.allclose(actual, self.expected_grad)
+
+    @pytest.mark.torch
+    def test_trainable_torch(self):
+        """Tests that coeffs passed to a sum are trainable with torch."""
+        import torch
+
+        dev = qml.device("default.qubit", wires=1)
+        qnode = qml.QNode(self.circuit, dev, interface="torch")
+        y, z = torch.tensor(1.1, requires_grad=True), torch.tensor(2.2, requires_grad=True)
+        qnode(y, z).backward()
+        actual = [y.grad, z.grad]
+        assert np.allclose(actual, self.expected_grad)
+
+    @pytest.mark.tf
+    def test_trainable_tf(self):
+        """Tests that coeffs passed to a sum are trainable with tf."""
+        import tensorflow as tf
+
+        dev = qml.device("default.qubit", wires=1)
+        qnode = qml.QNode(self.circuit, dev, interface="tensorflow")
+        y, z = tf.Variable(1.1, dtype=tf.float64), tf.Variable(2.2, dtype=tf.float64)
+        with tf.GradientTape() as tape:
+            res = qnode(y, z)
+        actual = tape.gradient(res, [y, z])
+        assert np.allclose(actual, self.expected_grad)
+
+    @pytest.mark.jax
+    def test_trainable_jax(self):
+        """Tests that coeffs passed to a sum are trainable with jax."""
+        import jax
+
+        dev = qml.device("default.qubit", wires=1)
+        qnode = qml.QNode(self.circuit, dev, interface="jax")
+        y, z = jax.numpy.array([1.1, 2.2])
+        actual = jax.grad(qnode, argnums=[0, 1])(y, z)
+        assert np.allclose(actual, self.expected_grad)
+
+
 class TestGetBatchSize:
     """Tests for the helper method ``_get_batch_size`` of ``QubitDevice``."""
 
diff --git a/tests/ops/op_math/test_sum.py b/tests/ops/op_math/test_sum.py
index fbd3eea45..102021e1c 100644
--- a/tests/ops/op_math/test_sum.py
+++ b/tests/ops/op_math/test_sum.py
@@ -1067,7 +1067,7 @@ class TestIntegration:
         dev = qml.device("default.qubit", wires=wires)
         sum_op = Sum(Prod(qml.RX(1.23, wires=0), qml.Identity(wires=1)), qml.Identity(wires=1))
 
-        @qml.qnode(dev)
+        @qml.qnode(dev, interface=None)
         def my_circ():
             qml.PauliX(0)
             return qml.expval(sum_op)
@@ -1079,7 +1079,7 @@ class TestIntegration:
         """Tests that the parameters of a Sum are considered trainable."""
         dev = qml.device("default.qubit", wires=2)
 
-        @qml.qnode(dev)
+        @qml.qnode(dev, interface=None)
         def circuit():
             return qml.expval(Sum(qml.RX(1.1, 0), qml.RY(qnp.array(2.2), 0)))