% Golden example 2 -- causal multi-head attention.
% Shows: leading axes as stack depth, a transpose that physically swaps the
% face, equal edge length on the contracted axis, and a Boolean mask drawn in
% a different grammar from the scores it gates.
% Layout is entirely by cursor: no absolute coordinate appears below.
%   ../scripts/build.sh mha-causal.tex
\documentclass[border=10pt]{standalone}
\usepackage[cjk]{supertensor}

\stsetrole{q}{stTeal}
\stsetrole{k}{stOrange}
\stsetrole{v}{stViolet}
\stsetrole{s}{stCoral}
\stsetrole{w}{stGray}

\stdim{T}{6}
\stdim{dh}{3}
\stdim{d}{9}       % d = h * d_h, h = 3

\begin{document}
\begin{tikzpicture}

% ============================================================ stage A row ===
\ststage{SA}{每头打分：沿 $d_h$ 收缩}
\strow{rowA}{T}
  \ststack[role=q, bracket=true]{Q}{}{T}{dh}{3}
  \stglyph{mA}{$\times$}
  % K^T: the face is physically swapped, not relabelled. Its height equals Q's
  % width -- that is the contracted axis d_h, drawn at one edge length.
  \ststack[role=k, bracket=true]{KT}{}{dh}{T}{3}
  \stglyph{eA}{$=$}
  \ststack[role=s]{S}{}{T}{T}{3}
\strowend
\stcaption{Q}{$\mathbf Q^{(i)}$}{$h\times T\times d_h$}
\stcaption{KT}{$\mathbf K^{(i)\top}$}{$h\times d_h\times T$}
\stcaption{S}{$\mathbf S^{(i)}$}{$h\times T\times T$}

% ============================================================ stage B row ===
\ststage{SB}{因果掩码与加权求和}
\strow{rowB}{T}
  % The mask is a Boolean support, not a magnitude: one flat level, exact
  % triangle, no stack -- it is shared by every head.
  \stface[role=w, pattern=data,
          data={300000,330000,333000,333300,333330,333333}]{M}{}{T}{T}
  \stlink{lB}{softmax}
  \ststack[role=s, pattern=causal]{A}{}{T}{T}{3}
  \stglyph{mB}{$\times$}
  \ststack[role=v, bracket=true]{V}{}{T}{dh}{3}
  \stglyph{eB}{$=$}
  \ststack[role=v]{O}{}{T}{dh}{3}
\strowend
\stcaption{M}{$\mathbf M$}{$T\times T$}
\stcaption{A}{$\mathbf A^{(i)}$}{$h\times T\times T$}
\stcaption{V}{$\mathbf V^{(i)}$}{$h\times T\times d_h$}
\stcaption{O}{$\mathbf O^{(i)}$}{$h\times T\times d_h$}

% ============================================================ stage C row ===
\ststage{SC}{沿 $d_h$ 拼接后投影}
\strow{rowC}{d}   % the d x d projection is the tallest object here
  % Concatenation reverses the split: gap=0pt makes the three h-shards of width
  % d_h tile a face of width d exactly, with no eyeballed offset.
  \stface[role=v]{C1}{}{T}{dh}
  \stface[role=v, gap=0pt]{C2}{}{T}{dh}
  \stface[role=v, gap=0pt]{C3}{}{T}{dh}
  \stglyph{mC}{$\times$}
  \stface[role=w]{WO}{}{d}{d}
  \stglyph{eC}{$=$}
  \stface[role=v, bracket=true]{Y}{}{T}{d}
\strowend
\stcaption{C2}{$[\,\mathbf O^{(1)}\mid\mathbf O^{(2)}\mid\mathbf O^{(3)}\,]$}{$T\times d$}
\stcaption{WO}{$\mathbf W_O$}{$d\times d$}
\stcaption{Y}{$\mathbf Y$}{$T\times d$}

% ================================================================= formula ==
% Placed last so it is centered on the figure that was actually drawn.
\sttopformula{F}{$\displaystyle
  \mathbf A^{(i)}=\mathrm{softmax}\!\left(
  \frac{\mathbf Q^{(i)}\mathbf K^{(i)\top}}{\sqrt{d_h}}+\mathbf M\right),\qquad
  \mathbf O^{(i)}=\mathbf A^{(i)}\mathbf V^{(i)}$}

% ============================================================== meaning box ==
\stbbox{all}
\stmeaningbox{mb}{16.8cm}{all}
  {$T$ 序列长度，$d_h$ 单头宽度，$h$ 头数（图中 $h=3$，即堆叠的三张面），
   $d=h\,d_h$；批轴 $B$ 省略}
  {$\mathbf S,\mathbf A$ 是分数与概率（行和为 $1$）；$\mathbf M\in\{0,-\infty\}^{T\times T}$
   是布尔支撑而非数值，被所有头共享，故只画一张、不堆叠；紫色一族标记 $\mathbf V\rightarrow\mathbf O\rightarrow\mathbf Y$ 同一数据流}
  {$(T\times d_h)(d_h\times T)\rightarrow(T\times T)$：$\mathbf K^{\top}$ 的面高即收缩维 $d_h$；
   拼接是切分的逆运算，$3$ 个 $d_h$ 恰好铺满 $d$}
\stsignature{因果多头注意力（掩码 + 拼接投影）}{mb}

\end{tikzpicture}
\end{document}