#import "@preview/fletcher:0.6.0" as fletcher: diagram, node, edge
#import fletcher.shapes: house, hexagon
#set text(font: "New Computer Modern")

#let blob(pos, label, tint: white, ..args) = node(
pos, align(center, label),
width: 28mm,
fill: tint.lighten(60%),
stroke: 1pt + tint.darken(20%),
corner-radius: 5pt,
..args,
)

#diagram(
spacing: 15pt,
edge-stroke: 1pt,
edge-corner-radius: 5pt,
mark-scale: 70%,

blob((0,1), [Add & Norm], tint: yellow, shape: "hexagon"),
edge(),
blob((0,2), [Multi-Head\ Attention], <mha>, tint: orange),
blob((0,4), [Input], <in>, shape: "house", angle: 30deg,
width: auto, tint: red),

edge(<in>, <mha>, "-n", mark-scale: 150%),
edge(<in>, "-|>", (rel: (+0.3,0), to: <mha>), corner: "|-|"),
edge(<in>, "-|>", (rel: (-0.3,0), to: <mha>), corner: "|-|"),

edge((0,4), "lu,uu,r", "--|>"),
edge((0,1), (0,0.35), "r", (1,3), "r,u", "-|>"),
edge((1,2), "d,rr,uu,l", "--|>"),

blob((2,0), [Softmax], tint: green),
edge("<|-"),
blob((2,1), [Add & Norm], tint: yellow, shape: "hexagon"),
edge(),
blob((2,2), [Feed\ Forward], tint: blue),
)