diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md
index 84cdf3160ca..7fb61fb68a1 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)
+
Contributors ✍️
This release contains contributions from (in alphabetical order):
diff --git a/pennylane/transforms/hamiltonian_expand.py b/pennylane/transforms/hamiltonian_expand.py
index 9e90ec80125..5baca494e41 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,
)
diff --git a/tests/transforms/test_hamiltonian_expand.py b/tests/transforms/test_hamiltonian_expand.py
index d4fa49cb6a1..e48549a9c6e 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 @@ 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."""