Matrices and shapes — Linear algebra for LLMs

A matrix is a grid of numbers

$A$ has 2 rows and 3 columns: an $m \times n$ matrix with $m = 2$ and $n = 3$ . We write $A \in \mathbb{R}^{2 \times 3}$ .

Write

A \in \mathbb{R}^{2 \times 3}

: 2 rows, 3 columns.

Rows go across, columns go down

Row 1 (blue) is $[1, 2, 3]$ . Column 2 (pink) is $[2, 5]$ . They cross at $A_{12} = 2$ (green), row first, then column.

Write

Row 1 =

[1, 2, 3]

. Column 2 =

[2, 5]

.

A_{12} = 2

.

✏️ Your turn

Read the entries

Write $A_{21}$ (math, counting from 1) and A[0, 2] (code, counting from 0).

Work it out on paper, then go to the next card to check.

Answer

$A_{21} = 4$ (row 2, column 1, in blue) and A[0, 2] = 3 (row 0, column 2 in code, in pink). Same grid, two counting conventions.

A stack of rows: one token per row

Read $X$ by rows and each row is one token’s vector. Row “cat” is everything the model holds about that token. Most layers act on each row separately.

Write

X

is

T \times d

. Row

i

= token

i

’s vector.

A stack of columns: one feature per column

Read $X$ by columns and each column is one feature across all the tokens: here feature 3 is $[2, -1, 0]$ . Normalization layers and some interpretability tools look along this direction.

Shapes you'll meet in a transformer

Matrix	Shape	One row is
$X$ , the hidden states	$T \times d$	one token
$E$ , the embedding table	$V \times d$	one vocabulary entry
$W_Q$ , a projection (code layout)	$d \times d_h$	one input feature’s weights
$QK^\top$ , the attention scores	$T \times T$	one query against every key

Write

X\!: T \times d

.

E\!: V \times d

.

W_Q\!: d \times d_h

. Scores:

T \times T

.

In PyTorch

A = torch.tensor([[1, 2, 3], [4, 5, 6]])
A.shape      # torch.Size([2, 3])
A[0]         # first row: tensor([1, 2, 3])
A[:, 1]      # second column: tensor([2, 5])
A[1, 0]      # one entry: tensor(4)

: means “every index along this axis”.

📌 Recap

Matrices and shapes in one card

An $m \times n$ matrix has $m$ rows and $n$ columns; $A_{ij}$ is row $i$ , column $j$ .
Math counts from 1 ( $A_{21}$ ); code counts from 0 (A[1, 0]).
Hidden states $X$ are $T \times d$ : each row is a token, each column a feature.

Write

The three lines above.