Hi @Tz_19 ,
You can take the figure and axes from the drawer and use them with Matplotlib methods such as annotate. Here’s an example:
import pennylane as qml
from matplotlib import pyplot as plt
def draw_attention_module():
drawer=qml.drawer.MPLDrawer(2, {i: i for i in range(2)})
drawer.box_gate(layer=0, wires=0, text='Rotation 1')
drawer.ctrl(layer=1, wires=[0, 1], control_values=[True, True]) # Add a second control wire and a second control value
drawer.box_gate(layer=1, wires=1, text='Rotation 2', box_options={"zorder": 3}) # Add box_options={"zorder": 3}
# use annotate
drawer.ax.annotate('annotate', xy=(0, 0), xytext=(2.1, 0))
qml.drawer.style._black_white()
draw_attention_module()
plt.show()
