Repository: PennyLaneAI/pennylane

Issue:
[BUG] QAOA layers incorrectly forbid `Sum` observables
### Expected behavior

A `Sum` observable should be treated just like a legacy hamiltonian object or a new `LinearCombination` object.

### Actual behavior

The `QAOA` layers manually raise an error during a validation step, forbidding `Sum` and other multi-term op math objects other than the legacy Hamiltonian and `LinearCombination`.

### Additional information

_No response_

### Source code

```shell
#imports
import pennylane as qml
from pennylane import numpy as np

from pennylane import qaoa

import cmath
import matplotlib.pyplot as plt

# cost and mixer hamiltonians
cost_H = qml.PauliZ(0) + qml.PauliZ(1) + qml.PauliZ(2)

mixer_H = qml.PauliX(0) + qml.PauliX(1) + qml.PauliX(2) 

# layers, qnode
wires = [0,1,2]
num_layers = 5

def qaoa_layer(gamma, alpha):
    qaoa.cost_layer(gamma, cost_H)
    qaoa.mixer_layer(alpha, mixer_H)
    
def circuit(params):
    for w in wires:
        qml.Hadamard(w)
    qml.layer(qaoa_layer, num_layers, params[0], params[1])
    
dev = qml.device("default.qubit", wires=wires)

@qml.qnode(dev)
def cost_function(params):
    circuit(params)
    return qml.expval(cost_H)

params = np.array([[0.5]*num_layers, [0.5]*num_layers], requires_grad = True)

cost_function(params)
```


### Tracebacks

```shell
---------------------------------------------------------------------------
ValueError                                Traceback (most recent call last)
Cell In[38], line 41
     38 params = np.array([[0.5]*num_layers, [0.5]*num_layers], requires_grad = True)
     40 for _ in range(steps):
---> 41     params = opt.step(cost_function, params)

File ~\anaconda3\Lib\site-packages\pennylane\optimize\gradient_descent.py:93, in GradientDescentOptimizer.step(self, objective_fn, grad_fn, *args, **kwargs)
     75 def step(self, objective_fn, *args, grad_fn=None, **kwargs):
     76     """Update trainable arguments with one step of the optimizer.
     77 
     78     Args:
   (...)
     90         If single arg is provided, list [array] is replaced by array.
     91     """
---> 93     g, _ = self.compute_grad(objective_fn, args, kwargs, grad_fn=grad_fn)
     94     new_args = self.apply_grad(g, args)
     96     # unwrap from list if one argument, cleaner return

File ~\anaconda3\Lib\site-packages\pennylane\optimize\gradient_descent.py:122, in GradientDescentOptimizer.compute_grad(objective_fn, args, kwargs, grad_fn)
    104 r"""Compute gradient of the objective function at the given point and return it along with
    105 the objective function forward pass (if available).
    106 
   (...)
    119     will not be evaluted and instead ``None`` will be returned.
    120 """
    121 g = get_gradient(objective_fn) if grad_fn is None else grad_fn
--> 122 grad = g(*args, **kwargs)
    123 forward = getattr(g, "forward", None)
    125 num_trainable_args = sum(getattr(arg, "requires_grad", False) for arg in args)

File ~\anaconda3\Lib\site-packages\pennylane\_grad.py:165, in grad.__call__(self, *args, **kwargs)
    162     self._forward = self._fun(*args, **kwargs)
    163     return ()
--> 165 grad_value, ans = grad_fn(*args, **kwargs)  # pylint: disable=not-callable
    166 self._forward = ans
    168 return grad_value

File ~\anaconda3\Lib\site-packages\autograd\wrap_util.py:20, in unary_to_nary.<locals>.nary_operator.<locals>.nary_f(*args, **kwargs)
     18 else:
     19     x = tuple(args[i] for i in argnum)
---> 20 return unary_operator(unary_f, x, *nary_op_args, **nary_op_kwargs)

File ~\anaconda3\Lib\site-packages\pennylane\_grad.py:183, in grad._grad_with_forward(fun, x)
    177 @staticmethod
    178 @unary_to_nary
    179 def _grad_with_forward(fun, x):
    180     """This function is a replica of ``autograd.grad``, with the only
    181     difference being that it returns both the gradient *and* the forward pass
    182     value."""
--> 183     vjp, ans = _make_vjp(fun, x)  # pylint: disable=redefined-outer-name
    185     if vspace(ans).size != 1:
    186         raise TypeError(
    187             "Grad only applies to real scalar-output functions. "
    188             "Try jacobian, elementwise_grad or holomorphic_grad."
    189         )

File ~\anaconda3\Lib\site-packages\autograd\core.py:10, in make_vjp(fun, x)
      8 def make_vjp(fun, x):
      9     start_node = VJPNode.new_root()
---> 10     end_value, end_node =  trace(start_node, fun, x)
     11     if end_node is None:
     12         def vjp(g): return vspace(x).zeros()

File ~\anaconda3\Lib\site-packages\autograd\tracer.py:10, in trace(start_node, fun, x)
      8 with trace_stack.new_trace() as t:
      9     start_box = new_box(x, t, start_node)
---> 10     end_box = fun(start_box)
     11     if isbox(end_box) and end_box._trace == start_box._trace:
     12         return end_box._value, end_box._node

File ~\anaconda3\Lib\site-packages\autograd\wrap_util.py:15, in unary_to_nary.<locals>.nary_operator.<locals>.nary_f.<locals>.unary_f(x)
     13 else:
     14     subargs = subvals(args, zip(argnum, x))
---> 15 return fun(*subargs, **kwargs)

File ~\anaconda3\Lib\site-packages\pennylane\workflow\qnode.py:1092, in QNode.__call__(self, *args, **kwargs)
   1089     override_shots = kwargs["shots"]
   1091 # construct the tape
-> 1092 self.construct(args, kwargs)
   1094 original_grad_fn = [self.gradient_fn, self.gradient_kwargs, self.device]
   1095 self._update_gradient_fn(shots=override_shots, tape=self._tape)

File ~\anaconda3\Lib\site-packages\pennylane\workflow\qnode.py:929, in QNode.construct(self, args, kwargs)
    926     self.interface = qml.math.get_interface(*args, *list(kwargs.values()))
    928 with qml.queuing.AnnotatedQueue() as q:
--> 929     self._qfunc_output = self.func(*args, **kwargs)
    931 self._tape = QuantumScript.from_queue(q, shots)
    933 params = self.tape.get_parameters(trainable_only=False)

Cell In[38], line 32, in cost_function(params)
     30 @qml.qnode(dev)
     31 def cost_function(params):
---> 32     circuit(params)
     33     return qml.expval(cost_H)

Cell In[38], line 26, in circuit(params)
     24 for w in wires:
     25     qml.Hadamard(w)
---> 26 qml.layer(qaoa_layer, num_layers, params[0], params[1])

File ~\anaconda3\Lib\site-packages\pennylane\templates\layer.py:213, in layer(template, depth, *args, **kwargs)
    211 for i in range(0, int(depth)):
    212     arg_params = [k[i] for k in args]
--> 213     template(*arg_params, **kwargs)

Cell In[38], line 20, in qaoa_layer(gamma, alpha)
     19 def qaoa_layer(gamma, alpha):
---> 20     qaoa.cost_layer(gamma, cost_H)
     21     qaoa.mixer_layer(alpha, mixer_H)

File ~\anaconda3\Lib\site-packages\pennylane\qaoa\layers.py:103, in cost_layer(gamma, hamiltonian)
     49 r"""Applies the QAOA cost layer corresponding to a cost Hamiltonian.
     50 
     51 For the cost Hamiltonian :math:`H_C`, this is defined as the following unitary:
   (...)
    100 
    101 """
    102 if not isinstance(hamiltonian, (qml.ops.Hamiltonian, qml.ops.LinearCombination)):
--> 103     raise ValueError(
    104         f"hamiltonian must be of type pennylane.Hamiltonian, got {type(hamiltonian).__name__}"
    105     )
    107 if not _diagonal_terms(hamiltonian):
    108     raise ValueError("hamiltonian must be written only in terms of PauliZ and Identity gates")

ValueError: hamiltonian must be of type pennylane.Hamiltonian, got Sum
```


### System information

```shell
master
```


### Existing GitHub issues

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

Discussion:
codecov[bot]: ## [Codecov](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5846?dropdown=coverage&src=pr&el=h1&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) Report
All modified and coverable lines are covered by tests :white_check_mark:
> Project coverage is 99.67%. Comparing base [(`c6a7cb0`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/c6a7cb0f15556cbf72f57513d0cef1609cefd1af?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`7114fdf`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/7114fdf2c51d10c744a49235ae7bf2197345de4c?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 2 commits behind head on master.


<details><summary>Additional details and impacted files</summary>


```diff
@@            Coverage Diff             @@
##           master    #5846      +/-   ##
==========================================
- Coverage   99.67%   99.67%   -0.01%     
==========================================
  Files         418      418              
  Lines       40080    39779     -301     
==========================================
- Hits        39951    39649     -302     
- Misses        129      130       +1     
```



</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5846?dropdown=coverage&src=pr&el=continue&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).   
:loudspeaker: Have feedback on the report? [Share it here](https://about.codecov.io/codecov-pr-comment-feedback/?utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

albi3ro: > LGTM! (what's an `OperationRecorder`???)

A variant of `QuantumTape` that has a slightly different str output that decomposes nested tapes immediately.  Most of the time isn't hurting anybody, but really shouldn't be using in places like this.

Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 25cefac72..5b2896486 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -326,6 +326,9 @@
 
 <h3>Bug fixes 🐛</h3>
 
+* `qml.qaoa.cost_layer` and `qml.qaoa.mixer_layer` can now be used with `Sum` operators.
+  [(#5846)](https://github.com/PennyLaneAI/pennylane/pull/5846)
+
 * Fixes a bug where `MottonenStatePreparation` produces wrong derivatives at special parameter values.
   [(#5774)](https://github.com/PennyLaneAI/pennylane/pull/5774)
 
diff --git a/pennylane/qaoa/layers.py b/pennylane/qaoa/layers.py
index 25712aa4d..4250806aa 100644
--- a/pennylane/qaoa/layers.py
+++ b/pennylane/qaoa/layers.py
@@ -28,21 +28,18 @@ def _diagonal_terms(hamiltonian):
     Returns:
         bool: ``True`` if all terms are products of diagonal Pauli gates, ``False`` otherwise
     """
-    val = True
 
-    for i in hamiltonian.ops:
-        if isinstance(i, Tensor):
-            obs = i.obs
-        elif isinstance(i, qml.ops.Prod):
-            obs = i.operands
+    for op in hamiltonian.terms()[1]:
+        if isinstance(op, Tensor):
+            obs = op.obs
+        elif isinstance(op, qml.ops.Prod):
+            obs = op.operands
         else:
-            obs = [i]
+            obs = [op]
         for j in obs:
             if j.name not in ("PauliZ", "Identity"):
-                val = False
-                break
-
-    return val
+                return False
+    return True
 
 
 def cost_layer(gamma, hamiltonian):
@@ -99,15 +96,10 @@ def cost_layer(gamma, hamiltonian):
         1: ──H───────────╰RZZ(1.00)─┤  <Z>
 
     """
-    if not isinstance(hamiltonian, (qml.ops.Hamiltonian, qml.ops.LinearCombination)):
-        raise ValueError(
-            f"hamiltonian must be of type pennylane.Hamiltonian, got {type(hamiltonian).__name__}"
-        )
-
+    op = qml.templates.ApproxTimeEvolution(hamiltonian, gamma, 1)
     if not _diagonal_terms(hamiltonian):
         raise ValueError("hamiltonian must be written only in terms of PauliZ and Identity gates")
-
-    qml.templates.ApproxTimeEvolution(hamiltonian, gamma, 1)
+    return op
 
 
 def mixer_layer(alpha, hamiltonian):
@@ -161,9 +153,4 @@ def mixer_layer(alpha, hamiltonian):
         1: ──H───────────╰RXX(1.00)─┤  <Z>
 
     """
-    if not isinstance(hamiltonian, (qml.ops.Hamiltonian, qml.ops.LinearCombination)):
-        raise ValueError(
-            f"hamiltonian must be of type pennylane.Hamiltonian, got {type(hamiltonian).__name__}"
-        )
-
-    qml.templates.ApproxTimeEvolution(hamiltonian, alpha, 1)
+    return qml.templates.ApproxTimeEvolution(hamiltonian, alpha, 1)


Test patch:
diff --git a/tests/test_qaoa.py b/tests/test_qaoa.py
index 6946214f1..8280fa79b 100644
--- a/tests/test_qaoa.py
+++ b/tests/test_qaoa.py
@@ -1160,6 +1160,10 @@ def make_mixer_layer_test_cases():
             qml.Hamiltonian([1, 1], [qml.PauliX(0), qml.PauliX(1)]),
             [qml.PauliRot(2, "X", wires=[0]), qml.PauliRot(2, "X", wires=[1])],
         ],
+        [
+            qml.X(0) + qml.X(1),
+            [qml.PauliRot(2, "X", wires=[0]), qml.PauliRot(2, "X", wires=[1])],
+        ],
         [
             qaoa.xy_mixer(Graph([(0, 1), (1, 2), (2, 0)])),
             [
@@ -1180,6 +1184,10 @@ def make_cost_layer_test_cases():
             qml.Hamiltonian([1, 1], [qml.PauliZ(0), qml.PauliZ(1)]),
             [qml.PauliRot(2, "Z", wires=[0]), qml.PauliRot(2, "Z", wires=[1])],
         ],
+        [
+            qml.Z(0) + qml.Z(1),
+            [qml.PauliRot(2, "Z", wires=[0]), qml.PauliRot(2, "Z", wires=[1])],
+        ],
         [
             qaoa.maxcut(Graph([(0, 1), (1, 2), (2, 0)]))[0],
             [
@@ -1199,7 +1207,9 @@ class TestLayers:
 
         hamiltonian = [[1, 1], [1, 1]]
 
-        with pytest.raises(ValueError, match=r"hamiltonian must be of type pennylane.Hamiltonian"):
+        with pytest.raises(
+            ValueError, match=r"hamiltonian must be a linear combination of pauli words"
+        ):
             qaoa.mixer_layer(0.1, hamiltonian)
 
     def test_cost_layer_errors(self):
@@ -1207,7 +1217,9 @@ class TestLayers:
 
         hamiltonian = [[1, 1], [1, 1]]
 
-        with pytest.raises(ValueError, match=r"hamiltonian must be of type pennylane.Hamiltonian"):
+        with pytest.raises(
+            ValueError, match=r"hamiltonian must be a linear combination of pauli words"
+        ):
             qaoa.cost_layer(0.1, hamiltonian)
 
         hamiltonian = qml.Hamiltonian([1, 1], [qml.PauliZ(0), qml.PauliX(1)])
@@ -1265,15 +1277,18 @@ class TestLayers:
 
         gamma = 1
 
-        with qml.tape.OperationRecorder() as rec:
-            qaoa.cost_layer(gamma, cost)
+        with qml.queuing.AnnotatedQueue() as q:
+            out = qaoa.cost_layer(gamma, cost)
 
-        rec = rec.expand()
+        expected = qml.ApproxTimeEvolution(cost, gamma, 1)
+        assert qml.equal(out, expected)
 
-        for i, j in zip(rec.operations, gates):
-            prep = [i.name, i.parameters, i.wires]
-            target = [j.name, j.parameters, j.wires]
-            assert prep == target
+        assert q.queue[0] is out
+        assert len(q) == 1
+        decomp = out.decomposition()
+
+        for i, j in zip(decomp, gates):
+            assert qml.equal(i, j)
 
     with qml.operation.disable_new_opmath_cm():
         cost_layer_test_cases_legacy = make_cost_layer_test_cases()