hamiltonian_expand support for Identity terms (PennyLaneAI#5414)

**Context:**

Since we now often switch back and forth from operations to pauli rep's
and back, we can end up with terms that are identities on no wires
`qml.I()`. Individual expectation values of such a term are forbidden
given may devices may not understand what to do with such a thing.
`qml.expval(qml.I())` raises a `NotImplementedError`.

But we still want to be able take expectation values of multi-term
observables with constant offsets.

**Description of the Change:**

This PR just updates `qml.transforms.hamiltonian_expand` to handle
multi-term observables with constant offsets.

**Benefits:**

`hamiltonian_expand` now works with multi-term observables with constant
offsets.

**Possible Drawbacks:**

This PR does not also update `sum_expand`, as it seems like that
transform first needs some organizational changes.

**Related GitHub Issues:**

[sc-55654]


**Note:** Pylint was complaining about the complexity of the function so
I broke it up a little bit. Makes more lines changed, but will probably
be easier to read now.

---------

Co-authored-by: Mudit Pandey <[email protected]>

doc/releases/changelog-dev.md

@@ -184,6 +184,9 @@
   [stim](https://github.com/quantumlib/Stim) `v1.13.0`.
   [(#5409)](https://github.com/PennyLaneAI/pennylane/pull/5409)
 
+* `qml.transforms.hamiltonian_expand` can now handle multi-term observables with a constant offset.
+  [(#5414)](https://github.com/PennyLaneAI/pennylane/pull/5414)
+
 <h4>Community contributions 🥳</h4>
 
 * Functions `measure_with_samples` and `sample_state` have been added to the new `qutrit_mixed` module found in

pennylane/devices/qubit/sampling.py

@@ -92,6 +92,13 @@ def _group_measurements(mps: List[Union[SampleMeasurement, ClassicalShadowMP, Sh
     return all_mp_groups, all_indices
 
 
+def _get_num_shots_for_expval_H(obs):
+    indices = obs.grouping_indices
+    if indices:
+        return len(indices)
+    return sum(int(not isinstance(o, qml.Identity)) for o in obs.terms()[1])
+
+
 # pylint: disable=no-member
 def get_num_shots_and_executions(tape: qml.tape.QuantumTape) -> Tuple[int, int]:
     """Get the total number of qpu executions and shots.
@@ -111,8 +118,7 @@ def get_num_shots_and_executions(tape: qml.tape.QuantumTape) -> Tuple[int, int]:
         if isinstance(group[0], ExpectationMP) and isinstance(
             group[0].obs, (qml.ops.Hamiltonian, qml.ops.LinearCombination)
         ):
-            indices = group[0].obs.grouping_indices
-            H_executions = len(indices) if indices else len(group[0].obs.ops)
+            H_executions = _get_num_shots_for_expval_H(group[0].obs)
             num_executions += H_executions
             if tape.shots:
                 num_shots += tape.shots.total_shots * H_executions

pennylane/operation.py

@@ -2506,13 +2506,12 @@ def prune(self):
 
         if len(self.non_identity_obs) == 0:
             # Return a single Identity as the tensor only contains Identities
-            obs = qml.Identity(self.wires[0])
-        elif len(self.non_identity_obs) == 1:
-            obs = self.non_identity_obs[0]
-        else:
-            obs = Tensor(*self.non_identity_obs)
-
-        return obs
+            return qml.Identity(self.wires[0]) if self.wires else qml.Identity()
+        return (
+            self.non_identity_obs[0]
+            if len(self.non_identity_obs) == 1
+            else Tensor(*self.non_identity_obs)
+        )
 
     def map_wires(self, wire_map: dict):
         """Returns a copy of the current tensor with its wires changed according to the given

pennylane/transforms/hamiltonian_expand.py

@@ -15,6 +15,7 @@
 Contains the hamiltonian expand tape transform
 """
 # pylint: disable=protected-access
+from functools import partial
 from typing import List, Sequence, Callable
 
 import pennylane as qml
@@ -24,6 +25,118 @@
 from pennylane.transforms import transform
 
 
+def grouping_processing_fn(res_groupings, coeff_groupings, batch_size, offset):
+    """Sums up results for the expectation value of a multi-term observable when grouping is involved.
+
+    Args:
+        res_groupings (ResultBatch): The results from executing the batch of tapes with grouped observables
+        coeff_groupings (List[TensorLike]): The coefficients in the same grouped structure as the results
+        batch_size (Optional[int]): The batch size of the tape and corresponding results
+        offset (TensorLike): A constant offset from the multi-term observable
+
+    Returns:
+        Result: The result of the expectation value for a multi-term observable
+    """
+    dot_products = []
+    for c_group, r_group in zip(coeff_groupings, res_groupings):
+        # pylint: disable=no-member
+        if isinstance(r_group, (tuple, list, qml.numpy.builtins.SequenceBox)):
+            r_group = qml.math.stack(r_group)
+        if qml.math.shape(r_group) == ():
+            r_group = qml.math.reshape(r_group, (1,))
+        if batch_size:
+            r_group = r_group.T
+
+        if len(c_group) == 1 and len(r_group) != 1:
+            dot_products.append(r_group * c_group)
+        else:
+            dot_products.append(qml.math.dot(r_group, c_group))
+    summed_dot_products = qml.math.sum(qml.math.stack(dot_products), axis=0)
+    interface = qml.math.get_deep_interface(res_groupings)
+    return qml.math.asarray(summed_dot_products + offset, like=interface)
+
+
+def _grouping_hamiltonian_expand(tape):
+    """Calculate the expectation value of a tape with a multi-term observable using the grouping
+    present on the observable.
+    """
+    hamiltonian = tape.measurements[0].obs
+    if hamiltonian.grouping_indices is None:
+        # explicitly selected grouping, but indices not yet computed
+        hamiltonian.compute_grouping()
+
+    coeff_groupings = []
+    obs_groupings = []
+    offset = 0
+    coeffs, obs = hamiltonian.terms()
+    for indices in hamiltonian.grouping_indices:
+        group_coeffs = []
+        obs_groupings.append([])
+        for i in indices:
+            if isinstance(obs[i], qml.Identity):
+                offset += coeffs[i]
+            else:
+                group_coeffs.append(coeffs[i])
+                obs_groupings[-1].append(obs[i])
+        coeff_groupings.append(qml.math.stack(group_coeffs))
+    # make one tape per grouping, measuring the
+    # observables in that grouping
+    tapes = []
+    for obs in obs_groupings:
+        new_tape = tape.__class__(tape.operations, (qml.expval(o) for o in obs), shots=tape.shots)
+
+        new_tape = new_tape.expand(stop_at=lambda obj: True)
+        tapes.append(new_tape)
+
+    batch_size = tape.batch_size
+
+    return tapes, partial(
+        grouping_processing_fn,
+        coeff_groupings=coeff_groupings,
+        batch_size=batch_size,
+        offset=offset,
+    )
+
+
+def naive_processing_fn(res, coeffs, offset):
+    """Sum up the results weighted by coefficients to get the expectation value of a multi-term observable.
+
+    Args:
+        res (ResultBatch): The result of executing a batch of tapes where each tape is a different term in the observable
+        coeffs (List(TensorLike)): The weights for each result in ``res``
+        offset (TensorLike): Any constant offset from the multi-term observable
+
+    Returns:
+        Result: the expectation value of the multi-term observable
+    """
+    dot_products = []
+    for c, r in zip(coeffs, res):
+        if qml.math.ndim(c) == 0 and qml.math.size(r) != 1:
+            dot_products.append(qml.math.squeeze(r) * c)
+        else:
+            dot_products.append(qml.math.dot(qml.math.squeeze(r), c))
+    summed_dot_products = qml.math.sum(qml.math.stack(dot_products), axis=0)
+    return qml.math.convert_like(summed_dot_products + offset, res[0])
+
+
+def _naive_hamiltonian_expand(tape):
+    """Calculate the expectation value of a multi-term observable using one tape per term."""
+    # make one tape per observable
+    hamiltonian = tape.measurements[0].obs
+    tapes = []
+    offset = 0
+    coeffs = []
+    for c, o in zip(*hamiltonian.terms()):
+        if isinstance(o, qml.Identity):
+            offset += c
+        else:
+            new_tape = tape.__class__(tape.operations, [qml.expval(o)], shots=tape.shots)
+            tapes.append(new_tape)
+            coeffs.append(c)
+
+    return tapes, partial(naive_processing_fn, coeffs=coeffs, offset=offset)
+
+
 @transform
 def hamiltonian_expand(tape: QuantumTape, group: bool = True) -> (Sequence[QuantumTape], Callable):
     r"""
@@ -113,93 +226,22 @@ def hamiltonian_expand(tape: QuantumTape, group: bool = True) -> (Sequence[Quant
 
     if (
         len(tape.measurements) != 1
-        or not isinstance(
-            hamiltonian := tape.measurements[0].obs,
-            (qml.ops.Hamiltonian, qml.ops.LinearCombination),
-        )
+        or not hasattr(tape.measurements[0].obs, "grouping_indices")
         or not isinstance(tape.measurements[0], ExpectationMP)
     ):
         raise ValueError(
-            "Passed tape must end in `qml.expval(H)`, where H is of type `qml.Hamiltonian`"
+            "Passed tape must end in `qml.expval(H)` where H can define grouping_indices"
         )
 
-    if qml.math.shape(hamiltonian.coeffs) == (0,) and qml.math.shape(hamiltonian.ops) == (0,):
+    hamiltonian = tape.measurements[0].obs
+    if len(hamiltonian.terms()[1]) == 0:
         raise ValueError(
             "The Hamiltonian in the tape has no terms defined - cannot perform the Hamiltonian expansion."
         )
 
-    # note: for backward passes of some frameworks
-    # it is crucial to use the hamiltonian.data attribute,
-    # and not hamiltonian.coeffs when recombining the results
-
     if group or hamiltonian.grouping_indices is not None:
-        if hamiltonian.grouping_indices is None:
-            # explicitly selected grouping, but indices not yet computed
-            hamiltonian.compute_grouping()
-
-        coeff_groupings = [
-            qml.math.stack([hamiltonian.data[i] for i in indices])
-            for indices in hamiltonian.grouping_indices
-        ]
-        obs_groupings = [
-            [hamiltonian.ops[i] for i in indices] for indices in hamiltonian.grouping_indices
-        ]
-
-        # make one tape per grouping, measuring the
-        # observables in that grouping
-        tapes = []
-        for obs in obs_groupings:
-            new_tape = tape.__class__(
-                tape.operations, (qml.expval(o) for o in obs), shots=tape.shots
-            )
-
-            new_tape = new_tape.expand(stop_at=lambda obj: True)
-            tapes.append(new_tape)
-
-        def processing_fn(res_groupings):
-            # pylint: disable=no-member
-            res_groupings = [
-                qml.math.stack(r) if isinstance(r, (tuple, qml.numpy.builtins.SequenceBox)) else r
-                for r in res_groupings
-            ]
-            res_groupings = [
-                qml.math.reshape(r, (1,)) if r.shape == () else r for r in res_groupings
-            ]
-            dot_products = []
-            for c_group, r_group in zip(coeff_groupings, res_groupings):
-                if tape.batch_size:
-                    r_group = r_group.T
-                if len(c_group) == 1 and len(r_group) != 1:
-                    dot_products.append(r_group * c_group)
-                else:
-                    dot_products.append(qml.math.dot(r_group, c_group))
-            summed_dot_products = qml.math.sum(qml.math.stack(dot_products), axis=0)
-
-            return qml.math.convert_like(summed_dot_products, res_groupings[0])
-
-        return tapes, processing_fn
-
-    coeffs = hamiltonian.data
-
-    # make one tape per observable
-    tapes = []
-    for o in hamiltonian.ops:
-        # pylint: disable=protected-access
-        new_tape = tape.__class__(tape.operations, [qml.expval(o)], shots=tape.shots)
-        tapes.append(new_tape)
-
-    # pylint: disable=function-redefined
-    def processing_fn(res):
-        dot_products = []
-        for c, r in zip(coeffs, res):
-            if qml.math.ndim(c) == 0 and qml.math.size(r) != 1:
-                dot_products.append(qml.math.squeeze(r) * c)
-            else:
-                dot_products.append(qml.math.dot(qml.math.squeeze(r), c))
-        summed_dot_products = qml.math.sum(qml.math.stack(dot_products), axis=0)
-        return qml.math.convert_like(summed_dot_products, res[0])
-
-    return tapes, processing_fn
+        return _grouping_hamiltonian_expand(tape)
+    return _naive_hamiltonian_expand(tape)
 
 
 # pylint: disable=too-many-branches, too-many-statements

tests/ops/op_math/test_sum.py

@@ -28,7 +28,6 @@
 from pennylane.operation import AnyWires, MatrixUndefinedError, Operator
 from pennylane.ops.op_math import Prod, Sum
 
-X, Y, Z = qml.PauliX, qml.PauliY, qml.PauliZ
 
 no_mat_ops = (
     qml.Barrier,

tests/test_vqe.py

@@ -337,7 +337,7 @@ def test_optimize_torch(self, dev_name, shots):
         exec_no_opt = tracker.totals["executions"]
 
         assert exec_opt == 5  # Number of groups in the Hamiltonian
-        assert exec_no_opt == 15
+        assert exec_no_opt == 14
 
         assert np.allclose(c1, c2, atol=1e-1)
 
@@ -383,7 +383,7 @@ def test_optimize_tf(self, shots):
         exec_no_opt = tracker.totals["executions"]
 
         assert exec_opt == 5  # Number of groups in the Hamiltonian
-        assert exec_no_opt == 15
+        assert exec_no_opt == 14
 
         assert np.allclose(c1, c2, atol=1e-1)
 
@@ -429,7 +429,7 @@ def test_optimize_autograd(self, shots):
         exec_no_opt = tracker.totals["executions"]
 
         assert exec_opt == 5  # Number of groups in the Hamiltonian
-        assert exec_no_opt == 15
+        assert exec_no_opt == 14
 
         assert np.allclose(c1, c2, atol=1e-1)
 

tests/transforms/test_hamiltonian_expand.py

@@ -262,6 +262,8 @@ def test_hamiltonian_dif_tensorflow(self):
 
         import tensorflow as tf
 
+        inner_dev = qml.device("default.qubit")
+
         H = qml.Hamiltonian(
             [-0.2, 0.5, 1], [qml.PauliX(1), qml.PauliZ(1) @ qml.PauliY(2), qml.PauliZ(0)]
         )
@@ -289,7 +291,7 @@ def test_hamiltonian_dif_tensorflow(self):
 
             tape = QuantumScript.from_queue(q)
             tapes, fn = hamiltonian_expand(tape)
-            res = fn(qml.execute(tapes, dev, qml.gradients.param_shift))
+            res = fn(qml.execute(tapes, inner_dev, qml.gradients.param_shift))
 
             assert np.isclose(res, output)
 
@@ -315,6 +317,50 @@ def circuit():
 
         assert res.shape == (3,)
 
+    def test_constant_offset_grouping(self):
+        """Test that hamiltonian_expand can handle a multi-term observable with a constant offset and grouping."""
+
+        H = 2.0 * qml.I() + 3 * qml.X(0) + 4 * qml.X(0) @ qml.Y(1) + qml.Z(0)
+        tape = qml.tape.QuantumScript([], [qml.expval(H)], shots=50)
+        batch, fn = qml.transforms.hamiltonian_expand(tape, group=True)
+
+        assert len(batch) == 2
+
+        tape_0 = qml.tape.QuantumScript([], [qml.expval(qml.Z(0))], shots=50)
+        tape_1 = qml.tape.QuantumScript(
+            [qml.RY(-np.pi / 2, 0), qml.RX(np.pi / 2, 1)],
+            [qml.expval(qml.Z(0)), qml.expval(qml.Z(0) @ qml.Z(1))],
+            shots=50,
+        )
+
+        assert qml.equal(batch[0], tape_0)
+        assert qml.equal(batch[1], tape_1)
+
+        dummy_res = (1.0, (1.0, 1.0))
+        processed_res = fn(dummy_res)
+        assert qml.math.allclose(processed_res, 10.0)
+
+    def test_constant_offset_no_grouping(self):
+        """Test that hamiltonian_expand can handle a multi-term observable with a constant offset and no grouping.."""
+
+        H = 2.0 * qml.I() + 3 * qml.X(0) + 4 * qml.X(0) @ qml.Y(1) + qml.Z(0)
+        tape = qml.tape.QuantumScript([], [qml.expval(H)], shots=50)
+        batch, fn = qml.transforms.hamiltonian_expand(tape, group=False)
+
+        assert len(batch) == 3
+
+        tape_0 = qml.tape.QuantumScript([], [qml.expval(qml.X(0))], shots=50)
+        tape_1 = qml.tape.QuantumScript([], [qml.expval(qml.X(0) @ qml.Y(1))], shots=50)
+        tape_2 = qml.tape.QuantumScript([], [qml.expval(qml.Z(0))], shots=50)
+
+        assert qml.equal(batch[0], tape_0)
+        assert qml.equal(batch[1], tape_1)
+        assert qml.equal(batch[2], tape_2)
+
+        dummy_res = (1.0, 1.0, 1.0)
+        processed_res = fn(dummy_res)
+        assert qml.math.allclose(processed_res, 10.0)
+
 
 with AnnotatedQueue() as s_tape1:
     qml.PauliX(0)