Matmul Flow ~ New Visualization Tool

AI by Hand ✍️ Academy

I built a new interactive tool called MatmulFlow to visualize the dimensions of matrix multiplications.

Matmul’s Dimensions are Confusing

If you’ve worked with matrix multiplication or matmul, you know the rule: if W’s shape is (m, n) and X’s shape is (n, p), then F = WX’s shape must be (m, p). The inner dimensions must match, the outer dimensions give you the result. It’s simple enough to state. But in practice, I struggled as a student in the past and I’ve watched my students struggle stumbling over this again and again. Which one is the inner dimension? Is it the columns of the first or the rows of the second?

And when you chain multiple multiplications together, it gets worse.

Consider W₅(W₄(W₃(W₂(W₁X)))). What happen if we increase W₃’s width by 40?

Consider ((((WX₁)X₂)X₃)X₄)X₅. What happen if X₂’s height increases from 150 to 300?

This is when my brain starts to hurt. But remember these questions for now, because we will revisit them using MatmulFlow later.

The problem with the standard notation, I think, is it is too abstract. You’re staring at letters — W, X, F, m, n, p — trying to keep track of which belongs to which. There’s nothing to see.

Just Shift the Second Matrix Up

Several years ago, I came up with a visual method that makes the shapes of matrix multiplication obvious at a glance. This method is extensible to complex chains of multiplication. No memorization needed. Here’s how it works.

Take your two matrices, W and X, and represent them as rectangles whose width and height correspond to the number of columns and rows. Now take the second matrix X and shift it up, placing it above and to the right of W.

Look at what happens. The right edge of W lines up with the bottom edge of X. That alignment is the inner dimension — the one that has to match for the multiplication to be valid. And the result F = WX simply fills in the remaining rectangle below X and to the right of W. It takes its height from W and its width from X.

The constraint that felt abstract — “the inner dimensions must match” — becomes geometric and obvious. If W’s right edge doesn’t line up with X’s bottom edge, the multiplication can’t happen. You don’t need to remember the rule. You can see it.

Chain of Matmuls

This visualization method extends naturally to chains of matrix multiplications. Each layer multiplies the output F of the previous layer by a new weight matrix W.

In MatmulFlow, stacking layers vertically means multiplying on the left. The output F₁ of the first layer feeds down to become X₂ for the next layer, giving F₂ = W₂ (W₁X₁).

The same visual logic applies in the other direction. Stacking layers horizontally means multiplying on the right. The output F₁ feeds across to become W₂ for the next layer, giving F₂ = (W₁X₁)X₂.

Let’s revisit W₅(W₄(W₃(W₂(W₁X)))) and (((W₁X₁)X₂)X₃)X₄)X₅. If we apply the same logic to visualize them, we can immediately see how they are different. The first expression stacks vertically, while the second expression stacks horizontally.

Resizing a Matrix

Back to our questions earlier, what happen if we increase W₃’s height by 40?

What happen if X₂’s height increases from 150 to 300?

Give MatmulFlow a try. Resize any matrix in the chain and watch the dimensions “flow” to other matrices in the chain in both directions.

Mastering Matrix Multiplication

If you want to commit yourself to mastering matrix multiplication, the paper version of my Matrix Multiplication Workbook was recently released on Amazon. Click this link to take a look.

If you are a member of the AI by Hand Academy, just a reminder that you already have access to the electronic version of this workbook.