diff --git a/doc/releases/changelog-dev.md b/doc/releases/changelog-dev.md index c9b1516689d..cbb10a19980 100644 --- a/doc/releases/changelog-dev.md +++ b/doc/releases/changelog-dev.md @@ -179,6 +179,9 @@

Improvements 🛠

+* Implemented the method `process_counts` in `ExpectationMp`. + [(#5241)](https://github.com/PennyLaneAI/pennylane/issues/5241) +

Faster gradients with VJPs and other performance improvements

* Adjoint device VJP's are now supported with `jax.jacobian`. `device_vjp=True` is @@ -683,4 +686,5 @@ Lee J. O'Riordan, Mudit Pandey, Alex Preciado, Matthew Silverman, -Jay Soni. +Jay Soni, +Tarun Kumar Allamsetty. diff --git a/pennylane/measurements/expval.py b/pennylane/measurements/expval.py index 6e8f1464932..e09b64a3028 100644 --- a/pennylane/measurements/expval.py +++ b/pennylane/measurements/expval.py @@ -136,3 +136,7 @@ def process_state(self, state: Sequence[complex], wire_order: Wires): prob = qml.probs(wires=self.wires).process_state(state=state, wire_order=wire_order) # In case of broadcasting, `prob` has two axes and this is a matrix-vector product return qml.math.dot(prob, eigvals) + + def process_counts(self, counts: dict, wire_order: Wires): + probs = qml.probs(wires=self.wires).process_counts(counts=counts, wire_order=wire_order) + return qml.math.dot(probs, self.eigvals()) diff --git a/tests/measurements/test_expval.py b/tests/measurements/test_expval.py index 21c36428de9..71e80e05c7d 100644 --- a/tests/measurements/test_expval.py +++ b/tests/measurements/test_expval.py @@ -318,3 +318,19 @@ def cost_circuit(params): energy_batched = cost_circuit(params) assert qml.math.allequal(energy_batched, energy) + + @pytest.mark.parametrize( + "wire, expected", + [ + (0, 0.0), + (1, 1.0), + ], + ) + def test_estimate_expectation_with_counts(self, wire, expected): + counts = {"000": 100, "100": 100} + + wire_order = qml.wires.Wires((0, 1, 2)) + + res = qml.expval(qml.Z(wire)).process_counts(counts=counts, wire_order=wire_order) + + assert np.allclose(res, expected)