|
| 1 | +/** |
| 2 | + * cppmega-mlx-w2t6: tiny SVG mini-framework for tensor-flow diagrams |
| 3 | + * inside HelpModal / ExplainModal. |
| 4 | + * |
| 5 | + * Each diagram is declarative JSX: a <Diagram> wrapper picks the |
| 6 | + * canvas size + dark-theme background, then primitives like |
| 7 | + * <Tensor>, <Op>, <Arrow>, <Residual>, <Group> position elements via |
| 8 | + * absolute x/y coords. The container automatically scales to its |
| 9 | + * container width while preserving the aspect ratio. |
| 10 | + * |
| 11 | + * Colour palette pinned to the rest of the canvas: |
| 12 | + * - background: rgba(15, 23, 42, 0.6) (#0f172a-ish) |
| 13 | + * - tensor box: fill #1e293b, stroke #22d3ee (cyan) |
| 14 | + * - op box: fill #312e81, stroke #818cf8 (indigo) |
| 15 | + * - arrow: #94a3b8 (slate-400) |
| 16 | + * - residual: #facc15 (yellow), dashed |
| 17 | + * - text: #f1f5f9 (slate-50) |
| 18 | + */ |
| 19 | + |
| 20 | +import { type ReactNode } from "react"; |
| 21 | + |
| 22 | + |
| 23 | +export const DIAG_THEME = { |
| 24 | + bg: "rgba(15, 23, 42, 0.6)", |
| 25 | + border: "rgba(255, 255, 255, 0.05)", |
| 26 | + tensorBg: "#1e293b", |
| 27 | + tensorFg: "#22d3ee", |
| 28 | + opBg: "#312e81", |
| 29 | + opFg: "#818cf8", |
| 30 | + arrow: "#94a3b8", |
| 31 | + residual: "#facc15", |
| 32 | + text: "#f1f5f9", |
| 33 | + textMuted: "#94a3b8", |
| 34 | +} as const; |
| 35 | + |
| 36 | + |
| 37 | +export interface DiagramProps { |
| 38 | + width?: number; |
| 39 | + height?: number; |
| 40 | + /** Diagram caption (rendered above the SVG). */ |
| 41 | + caption?: string; |
| 42 | + children: ReactNode; |
| 43 | +} |
| 44 | + |
| 45 | +/** Root SVG container. Scales to width via `width="100%"` on the |
| 46 | + * parent <svg>; coordinates inside are in the [0, width] × [0, height] |
| 47 | + * user-space frame. */ |
| 48 | +export function Diagram({ |
| 49 | + width = 480, height = 220, caption, children, |
| 50 | +}: DiagramProps): JSX.Element { |
| 51 | + return ( |
| 52 | + <div data-testid="tensor-diagram" style={{ |
| 53 | + background: DIAG_THEME.bg, |
| 54 | + border: `1px solid ${DIAG_THEME.border}`, |
| 55 | + borderRadius: 6, |
| 56 | + padding: 10, |
| 57 | + }}> |
| 58 | + {caption && ( |
| 59 | + <div data-testid="tensor-diagram-caption" style={{ |
| 60 | + color: DIAG_THEME.textMuted, |
| 61 | + fontSize: 11, |
| 62 | + marginBottom: 6, |
| 63 | + fontFamily: "ui-monospace, monospace", |
| 64 | + }}> |
| 65 | + {caption} |
| 66 | + </div> |
| 67 | + )} |
| 68 | + <svg viewBox={`0 0 ${width} ${height}`} |
| 69 | + width="100%" preserveAspectRatio="xMidYMid meet" |
| 70 | + style={{ display: "block" }}> |
| 71 | + <defs> |
| 72 | + <marker id="arrowhead" markerWidth="6" markerHeight="6" |
| 73 | + refX="6" refY="3" orient="auto" markerUnits="strokeWidth"> |
| 74 | + <path d="M0,0 L6,3 L0,6 Z" fill={DIAG_THEME.arrow} /> |
| 75 | + </marker> |
| 76 | + <marker id="arrowhead-residual" markerWidth="6" markerHeight="6" |
| 77 | + refX="6" refY="3" orient="auto" markerUnits="strokeWidth"> |
| 78 | + <path d="M0,0 L6,3 L0,6 Z" fill={DIAG_THEME.residual} /> |
| 79 | + </marker> |
| 80 | + </defs> |
| 81 | + {children} |
| 82 | + </svg> |
| 83 | + </div> |
| 84 | + ); |
| 85 | +} |
| 86 | + |
| 87 | + |
| 88 | +export interface TensorProps { |
| 89 | + x: number; |
| 90 | + y: number; |
| 91 | + w?: number; |
| 92 | + h?: number; |
| 93 | + label: string; |
| 94 | + /** Shape annotation rendered below the label, e.g. "(B, S, H)". */ |
| 95 | + shape?: string; |
| 96 | + testid?: string; |
| 97 | +} |
| 98 | + |
| 99 | +/** Rounded-rect tensor block. */ |
| 100 | +export function Tensor({ |
| 101 | + x, y, w = 90, h = 44, label, shape, testid, |
| 102 | +}: TensorProps): JSX.Element { |
| 103 | + return ( |
| 104 | + <g data-testid={testid}> |
| 105 | + <rect x={x} y={y} width={w} height={h} rx={6} |
| 106 | + fill={DIAG_THEME.tensorBg} |
| 107 | + stroke={DIAG_THEME.tensorFg} |
| 108 | + strokeWidth={1.5} /> |
| 109 | + <text x={x + w / 2} y={y + (shape ? h / 2 - 2 : h / 2 + 4)} |
| 110 | + textAnchor="middle" |
| 111 | + fill={DIAG_THEME.text} |
| 112 | + fontSize={12} fontWeight={600} |
| 113 | + fontFamily="system-ui, sans-serif"> |
| 114 | + {label} |
| 115 | + </text> |
| 116 | + {shape && ( |
| 117 | + <text x={x + w / 2} y={y + h / 2 + 12} |
| 118 | + textAnchor="middle" |
| 119 | + fill={DIAG_THEME.textMuted} |
| 120 | + fontSize={9} |
| 121 | + fontFamily="ui-monospace, monospace"> |
| 122 | + {shape} |
| 123 | + </text> |
| 124 | + )} |
| 125 | + </g> |
| 126 | + ); |
| 127 | +} |
| 128 | + |
| 129 | + |
| 130 | +export interface OpProps { |
| 131 | + x: number; |
| 132 | + y: number; |
| 133 | + w?: number; |
| 134 | + h?: number; |
| 135 | + label: string; |
| 136 | + testid?: string; |
| 137 | +} |
| 138 | + |
| 139 | +/** Operation node (matmul / softmax / silu / etc.). Hex-shaped. */ |
| 140 | +export function Op({ |
| 141 | + x, y, w = 70, h = 30, label, testid, |
| 142 | +}: OpProps): JSX.Element { |
| 143 | + return ( |
| 144 | + <g data-testid={testid}> |
| 145 | + <rect x={x} y={y} width={w} height={h} rx={14} |
| 146 | + fill={DIAG_THEME.opBg} |
| 147 | + stroke={DIAG_THEME.opFg} |
| 148 | + strokeWidth={1.2} /> |
| 149 | + <text x={x + w / 2} y={y + h / 2 + 4} |
| 150 | + textAnchor="middle" |
| 151 | + fill={DIAG_THEME.text} |
| 152 | + fontSize={11} |
| 153 | + fontFamily="system-ui, sans-serif"> |
| 154 | + {label} |
| 155 | + </text> |
| 156 | + </g> |
| 157 | + ); |
| 158 | +} |
| 159 | + |
| 160 | + |
| 161 | +export interface ArrowProps { |
| 162 | + x1: number; y1: number; x2: number; y2: number; |
| 163 | + /** Optional label rendered halfway along the arrow. */ |
| 164 | + label?: string; |
| 165 | + /** Bend along the y axis at this fraction of x1→x2. */ |
| 166 | + bendY?: number; |
| 167 | + testid?: string; |
| 168 | +} |
| 169 | + |
| 170 | +/** Straight or right-angle arrow between two anchors. */ |
| 171 | +export function Arrow({ |
| 172 | + x1, y1, x2, y2, label, bendY, testid, |
| 173 | +}: ArrowProps): JSX.Element { |
| 174 | + let path: string; |
| 175 | + let midX: number, midY: number; |
| 176 | + if (bendY !== undefined) { |
| 177 | + const mx = x1 + (x2 - x1) * (bendY ?? 0.5); |
| 178 | + path = `M ${x1} ${y1} L ${mx} ${y1} L ${mx} ${y2} L ${x2} ${y2}`; |
| 179 | + midX = mx; |
| 180 | + midY = (y1 + y2) / 2; |
| 181 | + } else { |
| 182 | + path = `M ${x1} ${y1} L ${x2} ${y2}`; |
| 183 | + midX = (x1 + x2) / 2; |
| 184 | + midY = (y1 + y2) / 2; |
| 185 | + } |
| 186 | + return ( |
| 187 | + <g data-testid={testid}> |
| 188 | + <path d={path} fill="none" |
| 189 | + stroke={DIAG_THEME.arrow} strokeWidth={1.5} |
| 190 | + markerEnd="url(#arrowhead)" /> |
| 191 | + {label && ( |
| 192 | + <text x={midX} y={midY - 5} textAnchor="middle" |
| 193 | + fill={DIAG_THEME.textMuted} |
| 194 | + fontSize={10} |
| 195 | + fontFamily="ui-monospace, monospace"> |
| 196 | + {label} |
| 197 | + </text> |
| 198 | + )} |
| 199 | + </g> |
| 200 | + ); |
| 201 | +} |
| 202 | + |
| 203 | + |
| 204 | +/** Dashed yellow residual line (skip connection). */ |
| 205 | +export function Residual({ |
| 206 | + x1, y1, x2, y2, label, bendY = 0.5, testid, |
| 207 | +}: ArrowProps): JSX.Element { |
| 208 | + const mx = x1 + (x2 - x1) * bendY; |
| 209 | + const path = |
| 210 | + `M ${x1} ${y1} L ${mx} ${y1} L ${mx} ${y2} L ${x2} ${y2}`; |
| 211 | + return ( |
| 212 | + <g data-testid={testid}> |
| 213 | + <path d={path} fill="none" |
| 214 | + stroke={DIAG_THEME.residual} strokeWidth={1.5} |
| 215 | + strokeDasharray="4 4" |
| 216 | + markerEnd="url(#arrowhead-residual)" /> |
| 217 | + {label && ( |
| 218 | + <text x={mx + 4} y={(y1 + y2) / 2} |
| 219 | + fill={DIAG_THEME.residual} |
| 220 | + fontSize={10} |
| 221 | + fontFamily="ui-monospace, monospace"> |
| 222 | + {label} |
| 223 | + </text> |
| 224 | + )} |
| 225 | + </g> |
| 226 | + ); |
| 227 | +} |
| 228 | + |
| 229 | + |
| 230 | +/** A logical group/box around a sub-region of the diagram (e.g. |
| 231 | + * multi-head attention). Renders a faint outlined rectangle with a |
| 232 | + * small caption in the top-left corner. */ |
| 233 | +export interface GroupProps { |
| 234 | + x: number; y: number; w: number; h: number; |
| 235 | + label?: string; |
| 236 | + testid?: string; |
| 237 | + children?: ReactNode; |
| 238 | +} |
| 239 | + |
| 240 | +export function Group({ |
| 241 | + x, y, w, h, label, testid, children, |
| 242 | +}: GroupProps): JSX.Element { |
| 243 | + return ( |
| 244 | + <g data-testid={testid}> |
| 245 | + <rect x={x} y={y} width={w} height={h} rx={8} |
| 246 | + fill="none" |
| 247 | + stroke="rgba(148, 163, 184, 0.3)" |
| 248 | + strokeWidth={1} |
| 249 | + strokeDasharray="2 4" /> |
| 250 | + {label && ( |
| 251 | + <text x={x + 6} y={y - 4} |
| 252 | + fill={DIAG_THEME.textMuted} |
| 253 | + fontSize={9} |
| 254 | + fontFamily="ui-monospace, monospace"> |
| 255 | + {label} |
| 256 | + </text> |
| 257 | + )} |
| 258 | + {children} |
| 259 | + </g> |
| 260 | + ); |
| 261 | +} |
0 commit comments