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."""