Repository: PennyLaneAI/pennylane
Pull request #5494 (2024-04-10 17:27:04)

Fix output dimensions bug in hamiltonian_expand
**Context:**
`hamiltonian_expand` reverses the order of dimensions for shot vectors and parameter broadcasting

**Description of the Change:**
Change a transpose operation to only transpose the last two dimensions if applicable.

**Related GitHub Issues:**
Fixes https://github.com/PennyLaneAI/pennylane/issues/5493
[sc-61037]

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/5494?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.66%. Comparing base [(`1089bb3`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/commit/1089bb3448081f4fde7bc3b4b792f1abda3555ff?dropdown=coverage&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI) to head [(`dba2dea`)](https://app.codecov.io/gh/PennyLaneAI/pennylane/pull/5494?dropdown=coverage&src=pr&el=desc&utm_medium=referral&utm_source=github&utm_content=comment&utm_campaign=pr+comments&utm_term=PennyLaneAI).
> Report is 3 commits behind head on master.


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


```diff
@@            Coverage Diff             @@
##           master    #5494      +/-   ##
==========================================
- Coverage   99.66%   99.66%   -0.01%     
==========================================
  Files         402      403       +1     
  Lines       37635    37389     -246     
==========================================
- Hits        37510    37263     -247     
- Misses        125      126       +1     
```



</details>

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


Patch:
diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 84cdf3160..7fb61fb68 100644
--- a/doc/releases/changelog-dev.md
+++ b/doc/releases/changelog-dev.md
@@ -447,6 +447,9 @@
 * `SampleMP`, `ExpectationMP`, `CountsMP`, `VarianceMP` constructed with ``eigvals`` can now properly process samples.
   [(#5463)](https://github.com/PennyLaneAI/pennylane/pull/5463)
 
+* Fixes a bug in `hamiltonian_expand` that produces incorrect output dimensions when shot vectors are combined with parameter broadcasting.
+  [(#5494)](https://github.com/PennyLaneAI/pennylane/pull/5494)
+
 <h3>Contributors ✍️</h3>
 
 This release contains contributions from (in alphabetical order):
diff --git a/pennylane/transforms/hamiltonian_expand.py b/pennylane/transforms/hamiltonian_expand.py
index 9e90ec801..5baca494e 100644
--- a/pennylane/transforms/hamiltonian_expand.py
+++ b/pennylane/transforms/hamiltonian_expand.py
@@ -44,13 +44,14 @@ def grouping_processing_fn(res_groupings, coeff_groupings, batch_size, offset):
             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 batch_size and batch_size > 1 and len(c_group) > 1:
+            r_group = qml.math.moveaxis(r_group, -1, -2)
 
         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)
@@ -88,12 +89,10 @@ def _grouping_hamiltonian_expand(tape):
         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,
+        batch_size=tape.batch_size,
         offset=offset,
     )
 


Test patch:
diff --git a/tests/transforms/test_hamiltonian_expand.py b/tests/transforms/test_hamiltonian_expand.py
index d4fa49cb6..e48549a9c 100644
--- a/tests/transforms/test_hamiltonian_expand.py
+++ b/tests/transforms/test_hamiltonian_expand.py
@@ -14,6 +14,8 @@
 """
 Unit tests for the ``hamiltonian_expand`` transform.
 """
+import functools
+
 import numpy as np
 import pytest
 
@@ -298,24 +300,65 @@ class TestHamiltonianExpand:
             g = gtape.gradient(res, var)
             assert np.allclose(list(g[0]) + list(g[1]), output2)
 
-    def test_processing_function_conditional_clause(self):
-        """Test the conditional logic for `len(c_group) == 1` and `len(r_group) != 1`
-        in the processing function returned by hamiltonian_expand, accessed when
-        using a shot vector and grouping if the terms don't commute with each other."""
-
-        dev_with_shot_vector = qml.device("default.qubit", shots=(10, 10, 10))
+    @pytest.mark.parametrize(
+        "H, expected",
+        [
+            # Contains only groups with single coefficients
+            (qml.Hamiltonian([1, 2.0], [qml.PauliZ(0), qml.PauliX(0)]), -1),
+            # Contains groups with multiple coefficients
+            (qml.Hamiltonian([1.0, 2.0, 3.0], [qml.X(0), qml.X(0) @ qml.X(1), qml.Z(0)]), -3),
+        ],
+    )
+    @pytest.mark.parametrize("grouping", [True, False])
+    def test_processing_function_shot_vectors(self, H, expected, grouping):
+        """Tests that the processing function works with shot vectors
+        and grouping with different number of coefficients in each group"""
 
-        H = qml.Hamiltonian([1, 2.0], [qml.PauliZ(0), qml.PauliX(0)])
-        H.compute_grouping()
+        dev_with_shot_vector = qml.device("default.qubit", shots=[(8000, 4)])
+        if grouping:
+            H.compute_grouping()
 
-        @qml.transforms.hamiltonian_expand
+        @functools.partial(qml.transforms.hamiltonian_expand, group=grouping)
         @qml.qnode(dev_with_shot_vector)
-        def circuit():
+        def circuit(inputs):
+            qml.RX(inputs, wires=0)
             return qml.expval(H)
 
-        res = circuit()
+        res = circuit(np.pi)
+        assert qml.math.shape(res) == (4,)
+        assert qml.math.allclose(res, np.ones((4,)) * expected, atol=0.1)
+
+    @pytest.mark.parametrize(
+        "H, expected",
+        [
+            # Contains only groups with single coefficients
+            (qml.Hamiltonian([1, 2.0], [qml.PauliZ(0), qml.PauliX(0)]), [1, 0, -1]),
+            # Contains groups with multiple coefficients
+            (
+                qml.Hamiltonian([1.0, 2.0, 3.0], [qml.X(0), qml.X(0) @ qml.X(1), qml.Z(0)]),
+                [3, 0, -3],
+            ),
+        ],
+    )
+    @pytest.mark.parametrize("grouping", [True, False])
+    def test_processing_function_shot_vectors_broadcasting(self, H, expected, grouping):
+        """Tests that the processing function works with shot vectors, parameter broadcasting,
+        and grouping with different number of coefficients in each group"""
+
+        dev_with_shot_vector = qml.device("default.qubit", shots=[(8000, 4)])
+
+        if grouping:
+            H.compute_grouping()
+
+        @functools.partial(qml.transforms.hamiltonian_expand, group=grouping)
+        @qml.qnode(dev_with_shot_vector)
+        def circuit(inputs):
+            qml.RX(inputs, wires=0)
+            return qml.expval(H)
 
-        assert res.shape == (3,)
+        res = circuit([0, np.pi / 2, np.pi])
+        assert qml.math.shape(res) == (4, 3)
+        assert qml.math.allclose(res, qml.math.stack([expected] * 4), atol=0.1)
 
     def test_constant_offset_grouping(self):
         """Test that hamiltonian_expand can handle a multi-term observable with a constant offset and grouping."""