Hello,

I am trying to reproduce the results of my own paper (https://arxiv.org/pdf/2302.12889.pdf) in which I used Yao.jl, in PennyLane. Here is the code snippet:

```
import pennylane as qml
import plotly.graph_objects as go
from IPython.display import display
from tqdm import tqdm
import time
# Quantum circuit and device setup
wires = 2
dev4 = qml.device("default.qubit", wires=wires, shots=1000)
def Q_Plot(cirq_0):
fig, ax = qml.draw_mpl(cirq_0, expansion_strategy='device')()
print(qml.draw(cirq_0, expansion_strategy='device')())
# plt.show()
fig.show()
@qml.qnode(dev4)
def ghz_circuit():
qml.Hadamard(wires=0)
qml.CNOT(wires=[0, 1])
# qml.CNOT(wires=[0, 2])
return qml.state()
# Initialize list to store results
all_results = []
# Create figure and histogram plots for the real and imaginary parts
fig_real = go.FigureWidget()
fig_real.add_bar(x=[f'Basis {i}' for i in range(2 ** wires)], y=[0] * (2 ** wires))
fig_real.update_layout(xaxis_title='Basis State', yaxis_title='Probability', title='Distribution of Real Part')
fig_imag = go.FigureWidget()
# Display the figures
display(fig_real)
# Generate N=1000 calls to the quantum circuit and plot their distribution in real time
N = 100
for _ in tqdm(range(N)):
result = ghz_circuit()
all_results.append(result)
# Separate the real and imaginary parts of the state vector
probs_real = abs(result.real) ** 2
# probs_imag = abs(result.imag) ** 2
# Update the histogram plots with new data
fig_real.data[0].y = probs_real
# fig_imag.data[0].y = probs_imag
# Update the plot layout
fig_real.update_layout(title=f'Distribution of Real Part ({len(all_results)}/{N})')
# fig_imag.update_layout(title=f'Distribution of Imaginary Part ({len(all_results)}/{N})')
# Pause to allow for interactive plotting
time.sleep(0.01)
Q_Plot(ghz_circuit)
```

This is the resulting plot:

And this is what I am trying to plot (from my paper):

Is there a better / built-in method to do this?

Thanks.