/tensor-formula-viz

name: tensor-formula-viz
description: Create or refine clean, shape-aware visualizations for tensor, matrix, and vector formulas or tensor code. Use for matrix-block diagrams, entry heatmaps, row/column shard stripes, stacked 3D/4D tensors, attention, tensor parallelism, broadcasting, reductions, contractions, and any explanation where tensor shapes or dimension meanings should be visually aligned with the computation.

Tensor Formula Viz

Turn a formula or tensor-code path into one dense, slide-ready figure with three zones:

1. **Top — formula:** clean mathematical definition, without shape underbraces. 2. **Middle — computation:** colored blocks, stripes, slices, operators, and communication. 3. **Bottom — meaning:** compact explanation of every shape dimension and non-obvious operation.

Workflow

1. Reduce the input to one primary computation path. Omit equivalent objectives, diagnostics, and secondary metrics unless requested. 2. Build a shape-and-semantics ledger: symbol/code name, semantic kind, dtype/domain, global and local shape, axis meanings, range when meaningful, shape-equivalence class, producer/consumer, and contracted/broadcast/reduced axes. 3. Choose the smallest visual grammar that exposes the mechanism. 4. Reserve non-overlapping stage lanes and object bounding boxes, then draw the three-zone layout. 5. Render, inspect, and correct both mathematical and visual alignment.

For code input, trace concrete `matmul`, `einsum`, `reshape/view`, `transpose/permute`, concat, broadcast, and collective operations. Preserve code variable names when useful; state any inferred shapes or conventions.

Choose the visual grammar

• **Entry color blocks:** use for masks, sparsity, causal structure, elementwise operations, or row/column roles inside a small matrix.
• **Row/column stripes:** use for sharding, device ownership, tensor parallelism, channel groups, or any operation whose key fact is an entire axis partition.
• **Stacked slices or head panels:** use for 3D/4D tensors such as \(B\times T\times d\) or \(B\times h\times T\times d_h\). Make depth/panels visibly encode the extra axes.
• **Arrows and device/collective nodes:** use for AllReduce, AllGather, ReduceScatter, transpose, reshape, split, concat, or data movement.

Combine grammars only when each adds information. Cell variation may give dense tensors visual texture, but known zeros stay unfilled and masks, diagonals, sparsity, and partitions must encode their exact structure.

Formula and shape rules

• Keep every displayed formula uncluttered. Never use shape underbraces in algebra; put shapes only under the corresponding blocks in the middle.
• Center the symbol directly below its block; center the shape on the next line.
• Explain dimension semantics at the bottom rather than repeating the shape list.
• Show the contracted dimension explicitly:

\[ (m\times k)(k\times n)\rightarrow(m\times n). \]

• For last-dimension linear layers, preserve leading axes:

\[ (B\times T\times d_{\mathrm{in}}) (d_{\mathrm{in}}\times d_{\mathrm{out}}) \rightarrow B\times T\times d_{\mathrm{out}}. \]

• Distinguish global shape from per-rank/local shape.
• Make transpose, broadcast, reduction, reshape, split, and concat axis order explicit.
• Treat colors as semantics: keep one color per tensor role, head, or TP rank throughout the figure.

Symbol semantics and discrete objects

• Classify every non-obvious symbol as value, score/logit, probability, index/coordinate, rank/order, count, ID, mask/support, permutation, or shape parameter. State its dtype/domain and what one entry means; shape alone is insufficient.
• Give one block one semantic object. Never merge a score, index list, and mask under a label such as \(M/S\); show each conversion with an explicit operator and arrow.
• For selection or routing, close the full chain: continuous scores \(\rightarrow\) discrete indices/IDs \(\rightarrow\) gather, scatter, mask, or route \(\rightarrow\) selected values. Distinguish `TopKValues` from `TopKIndices`; if both are used, show both outputs.
• Define index notation and range at first use, for example \(S_t=(s_{t,1},\ldots,s_{t,k})\), \(s_{t,r}\in\{0,\ldots,t\}\). Distinguish source-position axis \(s\) from selected-slot/rank axis \(r\), and state whether ordering, duplicates, padding, or variable cardinality matter.
• Show the address mapping once, such as \(G[b,t,r,:]=X[b,S[b,t,r],:]\). If a mask is also shown, state \(M[b,t,s]=\mathbf 1[s\in S_{b,t}]\); do not imply that the integer index tensor and Boolean mask are the same object.
• Visualize integer indices as ordered slots, position labels, or gather lines; visualize masks as binary structural support; visualize scores and values as magnitude color blocks. Do not render all three with the same heatmap grammar.
• When indices are shared across heads, ranks, or branches, draw the shared selector once and mark the broadcast/reuse axes explicitly.

Mandatory geometry invariants

• Map every matrix face \(a\times b\) to height \(a\) and width \(b\). For batched or stacked tensors, use the last two matrix axes for the face and leading axes only for depth or repeated panels.
• Derive one geometry ledger from the shape ledger. Assign each symbolic axis value one nominal physical edge length and reuse it across the entire figure. Equal shapes form an equivalence class and must have identical face width and height across stages, ranks, and semantic roles.
• Draw \(a\times a\) as a square. Draw a transpose by physically swapping face height and width; changing only the label is invalid.
• In \((m\times k)(k\times n)\), render both occurrences of the contracted \(k\) with the same physical edge length. Apply the same rule to `einsum` contractions and attention axes.
• Preserve partition geometry. Explicit shards tile their parent exactly along the split axis; equal shards have equal size; concatenation reverses the split. Size concatenated parts from their