diff --git a/src/braket/pennylane_plugin/translation.py b/src/braket/pennylane_plugin/translation.py index bd9eaebe..d477eaff 100644 --- a/src/braket/pennylane_plugin/translation.py +++ b/src/braket/pennylane_plugin/translation.py @@ -575,6 +575,8 @@ def translate_result_type( # noqa: C901 observable = measurement.obs if isinstance(measurement, qml.measurements.ProbabilityMP): + if observable and observable.diagonalizing_gates(): + raise qml.DeviceError("Probability result type not supported for observables") return Probability(targets) if isinstance(measurement, qml.measurements.StateMP): diff --git a/test/unit_tests/test_translation.py b/test/unit_tests/test_translation.py index ce11fdf4..ff15fbc2 100644 --- a/test/unit_tests/test_translation.py +++ b/test/unit_tests/test_translation.py @@ -829,6 +829,17 @@ def test_translate_result_type_unsupported_obs(): translate_result_type(tape.measurements[0], [0], frozenset()) +def test_translate_result_type_probs_observable(): + """Tests if a DeviceError is raised by translate_result_type for a Probability return type + with an observable attached""" + mp = qml.probs(op=qml.X(wires=0)) + + with pytest.raises( + qml.DeviceError, match="Probability result type not supported for observables" + ): + translate_result_type(mp, [0], frozenset()) + + def test_translate_result(): result_dict = _result_meta() result_dict["resultTypes"] = [