Repository: PennyLaneAI/pennylane

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.

Discussion:
github-actions[bot]: Hello. You may have forgotten to update the changelog!
Please edit [`doc/releases/changelog-dev.md`](/PennyLaneAI/pennylane/blob/master/doc/releases/changelog-dev.md) with:
* A one-to-two sentence description of the change. You may include a small working example for new features.
* A link back to this PR.
* Your name (or GitHub username) in the contributors section.
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4251?src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
> Merging [#4251](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4251?src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (e2c0215) into [master](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/c94ed82bed64c2f4b502e1f43edf28e2eda88f58?el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) (c94ed82) will **increase** coverage by `0.00%`.
> The diff coverage is `100.00%`.

```diff
@@           Coverage Diff           @@
##           master    #4251   +/-   ##
=======================================
  Coverage   99.77%   99.77%           
=======================================
  Files         352      352           
  Lines       31655    31667   +12     
=======================================
+ Hits        31584    31596   +12     
  Misses         71       71           
```


| [Impacted Files](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4251?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) | Coverage Δ | |
|---|---|---|
| [pennylane/devices/default\_qubit.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4251?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvZGVmYXVsdF9xdWJpdC5weQ==) | `99.70% <100.00%> (+<0.01%)` | :arrow_up: |
| [pennylane/devices/qubit/measure.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4251?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL2RldmljZXMvcXViaXQvbWVhc3VyZS5weQ==) | `100.00% <100.00%> (ø)` | |
| [pennylane/measurements/measurements.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4251?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL21lYXN1cmVtZW50cy9tZWFzdXJlbWVudHMucHk=) | `100.00% <100.00%> (ø)` | |
| [pennylane/ops/op\_math/sprod.py](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/4251?src=pr&el=tree&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI#diff-cGVubnlsYW5lL29wcy9vcF9tYXRoL3Nwcm9kLnB5) | `96.51% <100.00%> (+0.12%)` | :arrow_up: |


timmysilv: [sc-39422]

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)))