Transformer Block Viz

A stage-by-stage visualisation of the canonical transformer block: LayerNorm → MultiHead Attention → Add → LayerNorm → Feed-Forward → Add, with two curved residual bypass lines drawn alongside. A single cursor walks through the six stages; whichever residual currently spans the cursor brightens to full accent.

Installation

npx shadcn@latest add https://craftbits.dev/r/transformer-block-viz.json

Usage

import { TransformerBlockViz } from "@craft-bits/core";
 
<TransformerBlockViz />

Override stages to teach a post-norm block or any custom architecture:

<TransformerBlockViz
  stages={[
    { id: "mha",  label: "MultiHead Attention" },
    { id: "add1", label: "Add" },
    { id: "ln1",  label: "LayerNorm" },
    { id: "ffn",  label: "Feed-Forward" },
    { id: "add2", label: "Add" },
    { id: "ln2",  label: "LayerNorm" },
  ]}
  residuals={[
    { from: 0, to: 1 },
    { from: 2, to: 4 },
  ]}
/>

Drive playback from outside the component:

const [step, setStep] = useState(0);
 
<TransformerBlockViz
  currentStage={step}
  onCurrentStageChange={setStep}
  playing={false}
/>

Understanding the component

1. One cursor through the whole block. A single currentStage index drives the active highlight, the past / future styling, and which residual bypass is "lit" — there is no per-stage state machine.
2. Stages are cards; edges are arrows. Each stage renders as a bordered card with a step-number badge, label, and optional description. Edges between adjacent stages share a rail and an arrow-head glyph that points down (vertical) or right (horizontal).
3. Residuals are curved bypasses. Each entry in residuals is drawn as a cubic Bézier from the centre of from to the centre of to, bowing outward along the residual side. Any residual that currently spans the cursor brightens to full accent.
4. SPRINGS.smooth for the stage cursor. The active card scales to 1.02 and brightens via a smooth spring from @craft-bits/core/motion.
5. Reduced-motion fallback. With prefers-reduced-motion: reduce, the block renders fully lit on mount, autoplay pauses, and residual opacity transitions collapse to an instant swap.
6. Two orientations, same component. vertical (default — transformers are usually drawn top-to-bottom) and horizontal.

Props

Accessibility

• The root is role="figure" with aria-labelledby pointing at the "Transformer block" heading and aria-describedby at a visually-hidden aria-live="polite" summary.
• The summary announces the current stage and step counter whenever the cursor advances.
• Stages are rendered in a semantic <ol> with aria-current="step" on the active card.
• The play / pause button uses aria-pressed; the label flips between "Play block" and "Pause block".
• The scrubber is a native <input type="range"> with an aria-label so keyboard arrows nudge the cursor and screen readers narrate the value.
• Color is never the only signal — stage labels, descriptions, step-number badges, and the step counter are all textual.
• prefers-reduced-motion: reduce renders the block fully lit on mount, suppresses cursor / residual transitions, and disables autoplay.

Credits

• Extracted from: craftingattention (app/src/lessons/primitives/viz/TransformerBlockViz.tsx). Stripped the explore / predict / challenge mode strips, the pre-norm vs post-norm toggle, the Widget chrome with history-undo and bookmarks, the cell-tensor colour swatches, and the question-and-answer flow. Generalised to an arbitrary stages: TransformerBlockStage[] array with optional residuals, controlled / uncontrolled state pairs, and a horizontal layout option.