Measuring probabilities in the basis of an observable

[albi3ro]

Describe the bug

Pennylane now supports measuring probabilities in the basis on an observable.

For example,

>>> @qml.qnode(qml.device('default.qubit', wires=3))
... def circuit():
...     return qml.probs(op=qml.X(0))
>>> circuit()
array([0.5, 0.5])

In this case, qml.X(0)'s diagonalizing gates are applied prior to sampling.

Currently, the braket plugin ignores any diagonalizing gates the probability measurement might have:

>>> @qml.qnode(qml.device('braket.local.qubit', wires=3))
... def circuit():
...     return qml.probs(op=qml.X(0))
>>> circuit()
array([1., 0.])

This can lead to incorrect results with the hadamard gradient method:

>>> @qml.qnode(qml.device('braket.local.qubit', wires=3), diff_method="hadamard")
... def circuit(x):
...     qml.RX(x, 0)
...     return qml.probs(wires=0)
>>> qml.jacobian(circuit)(qml.numpy.array(0.5))
array([0., 0.])
>>> circuit.diff_method="parameter-shift"
>>> qml.jacobian(circuit)(qml.numpy.array(0.5))
array([-0.23971277,  0.23971277])
>>> print(qml.draw(qml.gradients.hadamard_grad(circuit))(qml.numpy.array(0.5)))
0: ──RX(0.50)─╭X────┤ ╭Probs[Z@Y]
1: ──H────────╰●──H─┤ ╰Probs[Z@Y]

To reproduce
A clear, step-by-step set of instructions to reproduce the bug.

The issue can be reproduced with the examples above.

Expected behavior
A clear and concise description of what you expected to happen.

I would either expect the device to apply the diagonalizing gates prior to measurement, or throw a NotImplementedError.

Screenshots or logs
If applicable, add screenshots or logs to help explain your problem.

System information
A description of your system. Please provide:

• Amazon Braket Python PennyLane Plugin version:1.30.1
• Amazon Braket Python SDK version:
• Amazon Braket Python Schemas version:
• Amazon Braket Python Default Simulator version:
• Python version: 3.11

Additional context
Add any other context about the problem here.