Hi @gidu,
Also, please how can I extract the probabilities for the different bases?
qml.probs returns probabilities of computational bases states. If you’d like to extract probabilities of states in another basis, then one approach is to apply rotations to the circuit followed by using qml.probs. The needed rotations can be obtained for an observable by calling the diagonalizing_gates method:
print(qml.Hadamard(wires=[0]).diagonalizing_gates())
[RY(-0.7853981633974483, wires=[0])]
This forum thread shares more details on how we obtain these rotations.
The code snippet seems to be using PennyLane well indeed. One thing to note, is that for n>3, qft would still only extract qml.probs(wires=[0, 1, 2]) (might be worth parametrizing the wires based on the value of n).