Repository: PennyLaneAI/pennylane
Pull request #6022 (2024-07-22 15:52:15)

Fix split_non_commuting bug for Hamiltonians with identity term
**Context:**
Described in https://github.com/PennyLaneAI/pennylane/issues/5924, but to summarize, `split_non_commuting` does not handle Identity terms with non-trainable coefficients well, leading to an unexpected `NonDifferentiableError`.

Although the coefficients of Identity terms (also known as `offset`) is assumed to be a float given the type-hinting of `_sum_terms()`, it seems that this is not the case. It is also not the case for `coeffs` which is assumed to be a list of floats but can actually be a list of tensors.

**Description of the Change:**
Fixes https://github.com/PennyLaneAI/pennylane/issues/5924 by setting `requires_grad` to `True` for offsets when autograd is in use. 

**Benefits:**
Fixes https://github.com/PennyLaneAI/pennylane/issues/5924

**Possible Drawbacks:**
May have unintended side effects, or perhaps we would like to coerce `offset` and `coeffs` to be the types in the type hint.

**Related GitHub Issues:**
[sc-67508]


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/6022?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.65%. Comparing base [(`3ad15d7`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/3ad15d7ee01b1c07b74d3d69f82c6c1699967076?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`22dc0bc`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/22dc0bc301224e150dadc673c879da901269fcf0?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).

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


```diff
@@            Coverage Diff             @@
##           master    #6022      +/-   ##
==========================================
- Coverage   99.66%   99.65%   -0.01%     
==========================================
  Files         427      427              
  Lines       41057    40764     -293     
==========================================
- Hits        40918    40624     -294     
- Misses        139      140       +1     
```

</details>

[:umbrella: View full report in Codecov by Sentry](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/6022?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).

austingmhuang: > Even if this fixes this particular problem, it will probably cause worse problems that will be extremely difficult to track down. Making something trainable can force us to do a lot more work later. Work the user explicitly didn't want. If someone marked something as trainable, but for some reason it starts being trainable, that could be really difficult to track down.
> 
> I'd prefer to find a different solution instead.

I'm not sure what the best fix for this particular problem is actually. 

How about this? We could add an additional check to see if the coeffs are trainable. Since the offset is just the coeff for the identity term, we can assume that the differentiability of the offset based on the differentiability of the coeffs. If coeff is trainable -> set `requires_grad` to `True` -- otherwise we set it to `False`.
astralcai: > Even if this fixes this particular problem, it will probably cause worse problems that will be extremely difficult to track down. Making something trainable can force us to do a lot more work later. Work the user explicitly didn't want. If someone marked something as trainable, but for some reason it starts being trainable, that could be really difficult to track down.
> 
> I'd prefer to find a different solution instead.

Unfortunately, after some investigation, it turned out that the reason why [this solution](https://github.com/PennyLaneAI/pennylane/pull/5945) worked was that it also caused things to become trainable, specifically via [this line](https://github.com/PennyLaneAI/pennylane/blob/master/pennylane/numpy/tensor.py#L170).
austingmhuang: > Unfortunately, after some investigation, it turned out that the reason why [this solution](https://github.com/PennyLaneAI/pennylane/pull/5945) worked was that it also caused things to become trainable, specifically via [this line](https://github.com/PennyLaneAI/pennylane/blob/master/pennylane/numpy/tensor.py#L170).

Is there a way to fix these two bugs without causing things to become trainable?
soranjh: > Since the offset is just the coeff for the identity term, we can assume the differentiability of the offset based on the differentiability of the coeffs. If coeff is trainable -> set requires_grad to True -- otherwise we set it to False.

This seems like a reasonable solution.
astralcai: I think there's a deeper issue here. The error is always raised from [here](https://github.com/PennyLaneAI/pennylane/blob/master/pennylane/numpy/tensor.py#L308), which forbids turning an untrainable tensor to an `ArrayBox`. When the code executes to [this line](https://github.com/PennyLaneAI/pennylane/blob/master/pennylane/transforms/split_non_commuting.py#L719), since there are other things in the circuit that are trainable, `sum_dot_products` is an `ArrayBox`, so executing `+` takes us [here](https://github.com/HIPS/autograd/blob/master/autograd/tracer.py#L31). This is a recursive function that takes values out of an `ArrayBox` [here](https://github.com/HIPS/autograd/blob/master/autograd/tracer.py#L37), performs the underlying operation [here](https://github.com/HIPS/autograd/blob/master/autograd/tracer.py#L44), and repackages the result into an `ArrayBox` [here](https://github.com/HIPS/autograd/blob/master/autograd/tracer.py#L46).

The issue is, when performing the underlying operation, we have on one hand the value that was taken out of the `ArrayBox` (which is now just a regular value), and on the other hand a tensor that was explicitly marked untrainable, the result is also an untrainable tensor. When autograd tries to repackage this result into an `ArrayBox`, it dispatches to [this function](https://github.com/PennyLaneAI/pennylane/blob/master/pennylane/numpy/tensor.py#L290), which now complains that it doesn't allow an untrainable tensor to be converted into an `ArrayBox`.
astralcai: > > Since the offset is just the coeff for the identity term, we can assume the differentiability of the offset based on the differentiability of the coeffs. If coeff is trainable -> set requires_grad to True -- otherwise we set it to False.
> 
> This seems like a reasonable solution to me. cc @albi3ro @astralcai

Maybe instead of assuming that all coeffs have the same trainability, we look at whether `summed_dot_products` is trainable, so something like:

```python
...
def _sum_terms(res: ResultBatch, coeffs: list[float], offset: float, shape: tuple) -> Result:
    ...
    summed_dot_products = qml.math.sum(qml.math.stack(dot_products), axis=0)
    if qml.math.get_interface(offset) == "autograd" and qml.math.requires_grad(summed_dot_products):
        offset = qml.math.array(offset)
    return summed_dot_products + offset
```

Patch:
diff --git a/pennylane/transforms/split_non_commuting.py b/pennylane/transforms/split_non_commuting.py
index 8c7c8b383..7082602ec 100644
--- a/pennylane/transforms/split_non_commuting.py
+++ b/pennylane/transforms/split_non_commuting.py
@@ -26,7 +26,7 @@ from pennylane.measurements import ExpectationMP, MeasurementProcess, Shots, Sta
 from pennylane.ops import Hamiltonian, LinearCombination, Prod, SProd, Sum
 from pennylane.tape import QuantumScript, QuantumScriptBatch
 from pennylane.transforms import transform
-from pennylane.typing import PostprocessingFn, Result, ResultBatch
+from pennylane.typing import PostprocessingFn, Result, ResultBatch, TensorLike, Union
 
 
 def null_postprocessing(results):
@@ -383,17 +383,17 @@ def _split_ham_with_grouping(tape: qml.tape.QuantumScript):
 
 def _split_using_qwc_grouping(
     tape: qml.tape.QuantumScript,
-    single_term_obs_mps: dict[MeasurementProcess, tuple[list[int], list[float]]],
-    offsets: list[float],
+    single_term_obs_mps: dict[MeasurementProcess, tuple[list[int], list[Union[float, TensorLike]]]],
+    offsets: list[TensorLike],
 ):
     """Split tapes using group_observables in the Pauli module.
 
     Args:
         tape (~qml.tape.QuantumScript): The tape to be split.
-        single_term_obs_mps (Dict[MeasurementProcess, Tuple[List[int], List[float]]]): A dictionary
+        single_term_obs_mps (Dict[MeasurementProcess, Tuple[List[int], List[TensorLike]]]): A dictionary
             of measurements of each unique single-term observable, mapped to the indices of the
             original measurements it belongs to, and its coefficients.
-        offsets (List[float]): Offsets associated with each original measurement in the tape.
+        offsets (List[TensorLike]): Offsets associated with each original measurement in the tape.
 
     """
 
@@ -449,17 +449,17 @@ def _split_using_qwc_grouping(
 
 def _split_using_wires_grouping(
     tape: qml.tape.QuantumScript,
-    single_term_obs_mps: dict[MeasurementProcess, tuple[list[int], list[float]]],
-    offsets: list[float],
+    single_term_obs_mps: dict[MeasurementProcess, tuple[list[int], list[Union[float, TensorLike]]]],
+    offsets: list[Union[float, TensorLike]],
 ):
     """Split tapes by grouping observables based on overlapping wires.
 
     Args:
         tape (~qml.tape.QuantumScript): The tape to be split.
-        single_term_obs_mps (Dict[MeasurementProcess, Tuple[List[int], List[float]]]): A dictionary
+        single_term_obs_mps (Dict[MeasurementProcess, Tuple[List[int], List[Union[float, TensorLike]]]]): A dictionary
             of measurements of each unique single-term observable, mapped to the indices of the
             original measurements it belongs to, and its coefficients.
-        offsets (List[float]): Offsets associated with each original measurement in the tape.
+        offsets (List[Union[float, TensorLike]]): Offsets associated with each original measurement in the tape.
 
     """
 
@@ -525,10 +525,10 @@ def _split_all_multi_term_obs_mps(tape: qml.tape.QuantumScript):
         tape (~qml.tape.QuantumScript): The tape with measurements to split.
 
     Returns:
-        single_term_obs_mps (Dict[MeasurementProcess, Tuple[List[int], List[float]]]): A
+        single_term_obs_mps (Dict[MeasurementProcess, Tuple[List[int], List[Union[float, TensorLike]]]]): A
             dictionary for measurements of each unique single-term observable, mapped to the
             indices of the original measurements it belongs to, and its coefficients.
-        offsets (List[float]): Offsets associated with each original measurement in the tape.
+        offsets (List[Union[float, TensorLike]]): Offsets associated with each original measurement in the tape.
 
     """
 
@@ -579,8 +579,8 @@ def _split_all_multi_term_obs_mps(tape: qml.tape.QuantumScript):
 
 def _processing_fn_no_grouping(
     res: ResultBatch,
-    single_term_obs_mps: dict[MeasurementProcess, tuple[list[int], list[float]]],
-    offsets: list[float],
+    single_term_obs_mps: dict[MeasurementProcess, tuple[list[int], list[Union[float, TensorLike]]]],
+    offsets: list[Union[float, TensorLike]],
     shots: Shots,
     batch_size: int,
 ):
@@ -589,10 +589,10 @@ def _processing_fn_no_grouping(
     Args:
         res (ResultBatch): The results from executing the tapes. Assumed to have a shape
             of (n_groups [,n_shots] [,n_mps] [,batch_size])
-        single_term_obs_mps (Dict[MeasurementProcess, Tuple[List[int], List[float]]]): A dictionary
+        single_term_obs_mps (Dict[MeasurementProcess, Tuple[List[int], List[Union[float, TensorLike]]]]): A dictionary
             of measurements of each unique single-term observable, mapped to the indices of the
             original measurements it belongs to, and its coefficients.
-        offsets (List[float]): Offsets associated with each original measurement in the tape.
+        offsets (List[Union[float, TensorLike]]): Offsets associated with each original measurement in the tape.
         shots (Shots): The shots settings of the original tape.
 
     """
@@ -631,8 +631,10 @@ def _processing_fn_no_grouping(
 
 def _processing_fn_with_grouping(
     res: ResultBatch,
-    single_term_obs_mps: dict[MeasurementProcess, tuple[list[int], list[float], int, int]],
-    offsets: list[float],
+    single_term_obs_mps: dict[
+        MeasurementProcess, tuple[list[int], list[Union[float, TensorLike]], int, int]
+    ],
+    offsets: list[Union[float, TensorLike]],
     group_sizes: list[int],
     shots: Shots,
     batch_size: int,
@@ -642,11 +644,11 @@ def _processing_fn_with_grouping(
     Args:
         res (ResultBatch): The results from executing the tapes. Assumed to have a shape
             of (n_groups [,n_shots] [,n_mps_in_group] [,batch_size])
-        single_term_obs_mps (Dict[MeasurementProcess, Tuple[List[int], List[float], int, int]]):
+        single_term_obs_mps (Dict[MeasurementProcess, Tuple[List[int], List[Union[float, TensorLike]], int, int]]):
             A dictionary of measurements of each unique single-term observable, mapped to the
             indices of the original measurements it belongs to, its coefficients, its group
             index, and the index of the measurement within the group.
-        offsets (List[float]): Offsets associated with each original measurement in the tape.
+        offsets (List[Union[float, TensorLike]]): Offsets associated with each original measurement in the tape.
         group_sizes (List[int]): The number of tapes in each group.
         shots (Shots): The shots setting of the original tape.
 
@@ -699,7 +701,12 @@ def _processing_fn_with_grouping(
     return tuple(res_for_each_mp)
 
 
-def _sum_terms(res: ResultBatch, coeffs: list[float], offset: float, shape: tuple) -> Result:
+def _sum_terms(
+    res: ResultBatch,
+    coeffs: list[Union[float, TensorLike]],
+    offset: Union[float, TensorLike],
+    shape: tuple,
+) -> Result:
     """Sum results from measurements of multiple terms in a multi-term observable."""
 
     # Trivially return the original result
@@ -714,7 +721,10 @@ def _sum_terms(res: ResultBatch, coeffs: list[float], offset: float, shape: tupl
         dot_products.append(qml.math.dot(qml.math.squeeze(r), c))
     if len(dot_products) == 0:
         return qml.math.ones(shape) * offset
+
     summed_dot_products = qml.math.sum(qml.math.stack(dot_products), axis=0)
+    if qml.math.get_interface(offset) == "autograd" and qml.math.requires_grad(summed_dot_products):
+        offset = qml.math.array(offset)
     return summed_dot_products + offset
 
 


Test patch:
diff --git a/tests/transforms/test_split_non_commuting.py b/tests/transforms/test_split_non_commuting.py
index ea71a4ff8..bc7853272 100644
--- a/tests/transforms/test_split_non_commuting.py
+++ b/tests/transforms/test_split_non_commuting.py
@@ -1056,7 +1056,7 @@ class TestDifferentiability:
             qml.RX(x, 0)
             c1 = qml.numpy.array(0.1, requires_grad=False)
             c2 = qml.numpy.array(0.2, requires_grad=False)
-            H = c1 * qml.Z(0) + c2 * qml.X(0)
+            H = c1 * qml.Z(0) + c2 * qml.X(0) + c2 * qml.I(0)
             return qml.expval(H)
 
         x = qml.numpy.array(0.5)