Fix output dimensions bug in hamiltonian_expand (#5494)

**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 #5493
[sc-61037]

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

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

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 @@ def test_hamiltonian_dif_tensorflow(self):
             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."""