Personalize this lesson
Adapt explanations and teaching visuals to your background and preferred voice.
Suppose three AI requests arrive with different prompt lengths and queue delays. You can use two numbers to predict one response time. How can the same small calculation handle all three requests, then keep a separate vector for every token in each prompt? The arithmetic stays familiar. The challenge is keeping track of which request, token, and feature each number describes.
The gradients lesson computed [4, 2] for a CodeAssist latency prediction: one loss slope for the prompt-length weight and one for the queue-delay weight. That ordered list is a vector. The gradient tells us how loss changes; the learning rate turns it into a weight update. Here, we'll follow those named values as one request becomes a batch and a batch becomes a token tensor.
Start with the vector you already computed
Before stacking anything, name the object in front of you. One number is a scalar, such as a loss of 1.0. An ordered list of numbers is a vector. A rectangular table of numbers is a matrix. The numbers don't carry their own meaning; axis labels and position conventions do.
Reuse the two-weight latency row:
| Vector | Meaning of position 0 | Meaning of position 1 | Values |
|---|---|---|---|
features | prompt length in hundreds of tokens | queue delay in seconds | [2, 1] |
weights | seconds per prompt block | seconds per second of queue delay | [2, 2] |
gradient | loss slope with respect to the prompt-length weight | loss slope with respect to the queue-delay weight | [4, 2] |
Each vector above has length 2, but they play distinct physical roles. Shape tells you the size of each axis, not the array's memory layout. Names and coordinate conventions tell you what the entries mean.
The NumPy lesson introduced .shape. A shape of (2,) denotes a single axis containing two items. That trailing comma is Python's syntax for a one-element tuple so it isn't confused with grouped parentheses. A NumPy scalar or zero-dimensional array has shape () because it has zero axes. An ordinary Python float doesn't have a .shape attribute.
Use NumPy arrays here because ordinary Python lists don't implement @. The learning-rate update itself is unchanged; NumPy lets us inspect the shape of each object as well as its values. All three inputs should have shape (2,), and subtracting 0.1 * gradient should give [1.6, 1.8]:
import numpy as np
features = np.array([2.0, 1.0])
weights = np.array([2.0, 2.0])
gradient = np.array([4.0, 2.0])
updated = weights - 0.1 * gradient
print("features", features.tolist(), "shape", features.shape)
print("weights", weights.tolist(), "shape", weights.shape)
print("gradient", gradient.tolist(), "shape", gradient.shape)
print("updated_weights", updated.tolist())features [2.0, 1.0] shape (2,)
weights [2.0, 2.0] shape (2,)
gradient [4.0, 2.0] shape (2,)
updated_weights [1.6, 1.8]An embedding is a learned vector representation. This vector holds two named features or two parameters; a token embedding later in the lesson has positions learned through training.
An update changes the weights. A dot product uses them to make the next prediction.
Let aligned features make one prediction
The latency model uses a dot product: multiply matching positions, then add. Let name the features, the weights, and ("y hat") the predicted latency. Subscripts 1 and 2 below identify the two coordinates; Python indexes those same positions with 0 and 1.
With the last chapter's starting weights [2, 2]:
That reproduces the 6-second prediction from the previous lesson. For the remaining latency examples, freeze a separate toy rule, w = [2, 1]. These chosen coefficients keep attention on array operations:
Prompt length contributes four seconds; queue delay contributes one. Two vectors of length two produce one scalar.
Check the difference between *, which keeps the two contributions, and @, which sums them. Both should agree with the hand calculation:
import numpy as np
features = np.array([2.0, 1.0])
weights = np.array([2.0, 1.0])
contributions = features * weights
prediction = features @ weights
print("contributions", contributions.tolist())
print("prediction_seconds", prediction)
print("output_shape", prediction.shape)
assert prediction == sum(a * b for a, b in zip(features, weights, strict=True))contributions [4.0, 1.0]
prediction_seconds 5.0
output_shape ()Element-wise multiplication stops at a contribution vector. The dot product adds those contributions into one number. For two real-valued 1-D arrays, NumPy's @ computes this inner product.[1]
Why can't you swap the feature order to [queue_delay, prompt_blocks] while keeping weights as [seconds per prompt block, seconds per second of queue delay]?
Answer
The lengths would still match, so the code would run, but positions would no longer describe the same quantities. You'd multiply queue delay by the prompt-length weight and prompt length by the queue-delay weight. Length checks catch structural errors; clear axis meaning catches semantic errors.
Now apply that same rule to several requests. We don't need a different weight vector for every row.
Stack requests into a matrix
Put three observed requests into a matrix: a rectangular array whose rows are requests and whose columns are features. The row axis lets one rule run across all requests.
Read conceptual shape [3, 2] as three requests by two features. Multiplying every row by the same weight vector returns one prediction per request:
The @ notation is the contract you'll see in library code. Here it means "dot each row with w." A length-two NumPy vector has shape (2,) even though we drew its values vertically in the equation. An explicit column matrix would have shape (2, 1) and would preserve a length-one output axis.
Predict the three results before running this cell: one short prompt with no queue gives 2; the other rows give 5 and 7:
import numpy as np
X = np.array([
[1.0, 0.0], # short prompt, empty queue
[2.0, 1.0], # medium prompt, one second of queue delay
[3.0, 1.0], # longer prompt, one second of queue delay
])
weights = np.array([2.0, 1.0])
predictions = X @ weights
print("X_shape", X.shape)
print("weights_shape", weights.shape)
print("predictions_shape", predictions.shape)
print("predicted_seconds", predictions.tolist())
assert (X @ weights[:, None]).shape == (3, 1)X_shape (3, 2)
weights_shape (2,)
predictions_shape (3,)
predicted_seconds [2.0, 5.0, 7.0]The one-request computation repeats across each row. Feature lengths must agree, and the row axis survives. The final assertion adds a singleton axis with [:, None] to highlight the distinction between a 1D vector and a 2D column matrix.
The three predictions use one shared rule. A second output needs another set of coefficients.
Use a matrix to make several outputs
The service may need more than one output. Suppose it also wants an escalation-priority score from the same two features. A weight vector produces one output; a weight matrix can produce both. The priority score uses arbitrary units and handpicked weights, so it carries no probability interpretation and has no trained decision boundary.
Use one column of W for each output:
The first column computes response seconds; the second computes priority. For x = [2, 1], run both dot products side by side:
The seconds column is the familiar [2, 1] rule. The priority column is a second dot product against the same features. The shape contract is:
For the whole three-row matrix:
This lesson stores examples as rows and outputs as columns, so it writes X @ W. Textbooks that store each example as a column may write W x instead. Both conventions work if axis meanings stay attached to the numbers.
![A 3 by 2 request matrix X multiplies a 2 by 2 weight matrix W to produce a 3 by 2 output Y. The highlighted row B in X is [2, 1], and the highlighted seconds column in W is [2, 1], so Y[B, seconds] is 2 times 2 plus 1 times 1, or 5. The shared feature axis of length 2 contracts; three request rows and two output columns remain. Matrix labels abbreviate prompt as p, queue as q, seconds as s, and priority as r.](/cdn/content-image/preparation/vectors-matrices-tensors/illustrations/_generated/matrix_multiply_shape_contract_dark.png?v=bdc54967ad29)
In the code, W[:, 0] selects the seconds column. Compare its predictions with the corresponding column of the full result:
import numpy as np
X = np.array([[1.0, 0.0], [2.0, 1.0], [3.0, 1.0]])
W = np.array([
[2.0, 0.2], # prompt contribution to [seconds, priority]
[1.0, 1.5], # queue contribution to [seconds, priority]
])
outputs = X @ W
print("contract", X.shape, "@", W.shape, "->", outputs.shape)
print("second_request", outputs[1].tolist())
print("all_seconds_predictions", outputs[:, 0].tolist())
np.testing.assert_allclose(outputs[:, 0], X @ W[:, 0])contract (3, 2) @ (2, 2) -> (3, 2)
second_request [5.0, 1.9]
all_seconds_predictions [2.0, 5.0, 7.0]Neural-net layers use the same idea at larger width: one learned matrix mixes input features into many output features for every row at once.
This rule also tells you when a proposed weight matrix cannot consume the available features.
Match inner dimensions before multiplying
Try reading the dimensions from the inside out. Each row of X has two features, but suppose W expects three:
X ends with a feature count of 2, while W starts with 3. There is no third feature to pair, so the product cannot be formed.
Run the invalid multiplication deliberately, then check a compatible matrix. The catch prints our own short diagnosis because NumPy's full error wording can vary by release:
import numpy as np
X = np.array([[1.0, 0.0], [2.0, 1.0], [3.0, 1.0]])
wrong_W = np.zeros((3, 2))
correct_W = np.array([[2.0, 0.2], [1.0, 1.5]])
try:
X @ wrong_W
except ValueError:
print("invalid", X.shape, "@", wrong_W.shape)
print("meeting_dimensions", X.shape[1], "and", wrong_W.shape[0], "do_not_match")
else:
raise AssertionError("expected incompatible input-feature dimensions")
print("valid_output_shape", (X @ correct_W).shape)invalid (3, 2) @ (3, 2)
meeting_dimensions 2 and 3 do_not_match
valid_output_shape (3, 2)When multiplication fails, name the axes first: request rows, input features, and output features. Then fix the object whose meaning is wrong.
Requests aren't the only extra axis. A prompt is a sequence of tokens, so a third axis must say which token each feature vector belongs to.
Add token and batch axes with tensors
In library code, a tensor means an array with any number of axes. NumPy calls its array type ndarray; PyTorch calls its type torch.Tensor. The practical array definition is the one used here and in the opening linear algebra chapter of Deep Learning.[2]
A scalar has zero axes, a vector one, and a matrix two. The number of axes is often called tensor rank, available as .ndim in NumPy. Matrix rank instead counts independent directions, which is a later question.
So far, we represented each request with two summary features: prompt length and queue delay. To represent its text, switch to one vector per token. These token coordinates aren't the latency features, even though we'll also give them length two. The NumPy lesson used the same (B, T, D) convention on log snippets:
| Axis | Symbol | Meaning here |
|---|---|---|
| batch | B | several prompts processed together |
| sequence | T | token positions within each prompt |
| feature | D | numbers describing one token at one layer |
Use tiny handcrafted token features so every number stays readable. This shape exercise uses a hand-built table:
| Token | Feature 0 | Feature 1 |
|---|---|---|
| test-like | 1.0 | 0.0 |
| fixture-like | 0.5 | 0.5 |
| import-like | 0.0 | 1.0 |
Two prompts, each containing three token-feature vectors, therefore have shape [2, 3, 2].
![Tensor H has axes batch, token, and feature, with shape [2,3,2]. Two 3 by 2 prompt matrices show all six token vectors. In prompt 1, the cell at token 2 and feature 0 is highlighted at 1.0. Following indices b=1, t=2, and d=0 reduces [2,3,2] to [3,2], then [2], then scalar 1.0.](/cdn/content-image/preparation/vectors-matrices-tensors/illustrations/_generated/vector_matrix_tensor_map_dark.png?v=c32c040fff91)
Read that shape as B = 2 prompts, T = 3 token positions, and D = 2 features per token. Indexing is zero-based, so H[1, 2, 0] selects the first feature of the third token in the second prompt. A colon keeps all entries along its axis: H[:, :, 0] therefore keeps batch and sequence but fixes the feature. Run both selections:
import numpy as np
H = np.array([
[[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]], # "test fixture import"
[[0.0, 1.0], [0.5, 0.5], [1.0, 0.0]], # "import fixture test"
])
print("tensor_shape", H.shape, "axes", H.ndim)
print("second_prompt_shape", H[1].shape)
print("third_token", H[1, 2].tolist(), "shape", H[1, 2].shape)
print("H[1,2,0]", H[1, 2, 0])
print("feature_zero_slice", H[:, :, 0].shape)
assert H.size == 2 * 3 * 2tensor_shape (2, 3, 2) axes 3
second_prompt_shape (3, 2)
third_token [1.0, 0.0] shape (2,)
H[1,2,0] 1.0
feature_zero_slice (2, 3)The last axis still contains features. Extra axes only organize which prompt and token each vector belongs to. A learned projection should change that last axis and leave the others alone.
Project each token without mixing positions
In this lesson, a projection is a matrix transformation of each token's features. Neural networks learn its coefficients; we'll choose small ones to inspect the arithmetic. Let matrix P map two input features into three output features:
Predict the output before multiplying: batch and sequence should remain intact, while the token-feature size changes from 2 to 3.
import numpy as np
H = np.array([
[[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]],
[[0.0, 1.0], [0.5, 0.5], [1.0, 0.0]],
])
P = np.array([
[1.0, 0.0, 0.5],
[0.0, 1.0, 0.5],
])
projected = H @ P
by_hand = [[[sum(token[d] * P[d, j] for d in range(2))
for j in range(3)] for token in prompt] for prompt in H]
print("contract", H.shape, "@", P.shape, "->", projected.shape)
print("first_prompt_first_token", projected[0, 0].tolist())
print("middle_token_shared", np.array_equal(projected[0, 1], projected[1, 1]))
np.testing.assert_allclose(projected, by_hand)contract (2, 3, 2) @ (2, 3) -> (2, 3, 3)
first_prompt_first_token [1.0, 0.0, 0.5]
middle_token_shared TrueRead the inner calculation as "multiply input coordinate d by its coefficient for output coordinate j, then sum over d." The nested Python loops spell out that calculation for every output feature, token, and prompt. np.testing.assert_allclose checks that the NumPy result matches this explicit version, allowing for small floating-point differences.
For a shared projection with shape [D, output_width], each token is transformed independently. Prompts don't mix, and token positions don't mix. More generally, @ multiplies the last two axes and broadcasts compatible leading batch axes.[1] Here P has no batch axis, so the same matrix serves both prompts.
This also explains an important limit: making the last dimension wider doesn't create new information. Each output is still a weighted combination of the two inputs. A later layer must do more than change width if it needs interactions between token positions.
A token tensor has shape [4, 128, 768] and a projection has shape [768, 256]. What output shape should you expect?
Answer
[4, 128, 256]. The meeting feature dimension 768 contracts, while batch and sequence axes remain intact.
Multiplication usually fails loudly when inner sizes disagree. Addition can stay quiet while still attaching a number to the wrong axis. Treat broadcasting as an axis-semantics problem: an arithmetically valid result can still attach values to the wrong feature.
Treat broadcasting as an axis decision
Broadcasting lets a smaller array participate in element-wise operations against a larger one. First choose what should repeat. An output bias [seconds_offset, priority_offset] should repeat across every request row:
NumPy compares dimensions from the right: equal sizes match, a size of 1 can expand, and missing leading dimensions act like size 1. The rule is useful, but it can also hide an axis mistake.[3]
Write both additions explicitly, then predict the second request. Feature bias should produce [5.1, 1.8]; row offsets should produce [4.9, 1.8] while still looking valid:
import numpy as np
outputs = np.array([
[2.0, 0.2],
[5.0, 1.9],
[7.0, 2.1],
]) # columns: [seconds, priority]
feature_bias = np.array([0.1, -0.1]) # one value per output feature
row_offsets = np.array([[0.1], [-0.1], [0.0]]) # one value per request
corrected = outputs + feature_bias
wrong_but_valid = outputs + row_offsets
print("correct_second_request", corrected[1].round(1).tolist())
print("wrong_second_request", wrong_but_valid[1].round(1).tolist())
assert corrected.shape == wrong_but_valid.shape == (3, 2)
assert feature_bias.shape == (outputs.shape[-1],)
assert not np.allclose(corrected, wrong_but_valid)correct_second_request [5.1, 1.8]
wrong_second_request [4.9, 1.8]Both results have three rows and two columns, but only one uses the intended axis. Checking only the result shape misses this bug. Requiring a one-dimensional bias with length equal to the output width catches the (3, 1) input. You still need named columns and a known-value example to catch two feature biases supplied in the wrong order.
The same dot product that scored one request can compare token vectors with each other. The result changes from one score to a grid of scores.
Read a score grid as repeated dot products
For three token-feature vectors in one prompt, multiply the token matrix by its transpose, which swaps rows and columns. The transpose turns each token row into a column so a row-column dot product compares two tokens. Predict the result shape first:
Every entry in the output is one dot product between two token positions. Rows name the first token; columns name the second.
| test-like | fixture-like | import-like | |
|---|---|---|---|
| test-like | 1.0 | 0.5 | 0.0 |
| fixture-like | 0.5 | 0.5 | 0.5 |
| import-like | 0.0 | 0.5 | 1.0 |
These are raw alignment scores. Cosine similarity requires dividing by both vector lengths. Notice the middle token's self-score: 0.5² + 0.5² = 0.5, while a unit-normalized vector has self-score 1. A dot product depends on vector length as well as direction.
For a real vector , its Euclidean length, written , is the square root of the sum of its squared coordinates. For nonzero vectors, cosine similarity is .[2] Dividing each vector by its length makes its self-score 1. A zero vector has no direction, so cosine similarity needs an explicit policy for that case.
You can draw a two-coordinate vector as an arrow from (0, 0) to its coordinates. Cosine similarity measures the angle between two such arrows: 1 for the same direction, 0 for perpendicular directions, and -1 for opposite directions. Scaling an arrow by a positive number changes its length but not that angle.
![Two equal-scale coordinate plots compare a=[0.5,0.5] and 2a=[1,1] with unit vector u=[1,0]. Both vectors point along the same 45-degree direction. The projected green segment and dot product grow from 0.50 to 1.00 as the vector doubles in length, while cosine similarity stays 1 divided by square root of 2, about 0.707.](/cdn/content-image/preparation/vectors-matrices-tensors/illustrations/_generated/dot_product_alignment_dark.png?v=14799e06a2b1)
Compute both grids and compare their diagonals. keepdims=True gives lengths shape [3, 1], keeping each length attached to its row during normalization:
import numpy as np
tokens = np.array([
[1.0, 0.0], # test-like
[0.5, 0.5], # fixture-like
[0.0, 1.0], # import-like
])
scores = tokens @ tokens.T # .T swaps the two axes of this matrix
lengths = np.linalg.norm(tokens, axis=1, keepdims=True)
if np.any(lengths == 0):
raise ValueError("cosine similarity needs nonzero token vectors")
unit_tokens = tokens / lengths
cosines = unit_tokens @ unit_tokens.T
print("contract", tokens.shape, "@", tokens.T.shape, "->", scores.shape)
print("raw_scores", scores.tolist())
print("cosine_diagonal", np.diag(cosines).round(6).tolist())
np.testing.assert_allclose(np.diag(cosines), 1.0)contract (3, 2) @ (2, 3) -> (3, 3)
raw_scores [[1.0, 0.5, 0.0], [0.5, 0.5, 0.5], [0.0, 0.5, 1.0]]
cosine_diagonal [1.0, 1.0, 1.0]Later, attention will create separate query and key vectors, scale the Q @ K.T score grid, and convert each row into weights that sum to one.[4] That expression is for one 2-D pair of matrices. For a batch shaped [B, T, D], use H @ H.swapaxes(-1, -2) to get [B, T, T]. Plain H.T reverses all three axes and would put batch size in the wrong place.
For the next exercise, use only the first two tokens of each prompt. All three axes now have length 2. Reversing every axis leaves shape (2, 2, 2) unchanged, and a legal multiply produces the expected output shape with the wrong values. Compare against each prompt's own 2-D score grid:
import numpy as np
H = np.array([
[[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]],
[[0.0, 1.0], [0.5, 0.5], [1.0, 0.0]],
])
H = H[:, :2, :] # two prompts, two token positions, two features
correct = H @ H.swapaxes(-1, -2)
wrong = H @ H.T
expected = np.stack([prompt @ prompt.T for prompt in H])
np.testing.assert_allclose(correct, expected)
assert wrong.shape == correct.shape == (2, 2, 2)
assert not np.allclose(wrong, expected)
print("Original and H.T shapes:", H.shape, H.T.shape)
print("Both output shapes:", correct.shape)
print("Prompt 0 self-scores, correct:", np.diag(correct[0]).tolist())
print("Prompt 0 self-scores, wrong: ", np.diag(wrong[0]).tolist())Original and H.T shapes: (2, 2, 2) (2, 2, 2)
Both output shapes: (2, 2, 2)
Prompt 0 self-scores, correct: [1.0, 0.5]
Prompt 0 self-scores, wrong: [1.0, 0.25]Each self-score should equal that token vector's squared length. The wrong diagonal violates that known-value check. Equal axis sizes don't make the axes interchangeable; they make semantic mistakes harder to detect.
If Q has shape [5, 8] and K has shape [5, 8], why does Q @ K.T have shape [5, 5]?
Answer
Transposing K turns [5, 8] into [8, 5]. The feature dimensions 8 meet, and the surviving dimensions describe five query positions by five key positions.
Build it: trace one tiny prompt projection
Now combine the pieces in a function that accepts a batch of token sequences and a projection matrix. Convert the inputs to numeric arrays first, reject the wrong number of axes, then check feature width. np.asarray(..., dtype=float) also rejects ragged lists such as [[1, 0], [0.5]]; inspecting only the first row wouldn't catch that missing coordinate.
Use the same two-to-three projection as before. After projecting, average over token positions to produce one summary vector per prompt. NumPy's mean(axis=1) removes the sequence axis and retains batch and features.[5]
import numpy as np
def project_prompts(hidden, projection):
try:
hidden = np.asarray(hidden, dtype=float)
projection = np.asarray(projection, dtype=float)
except (TypeError, ValueError) as error:
raise ValueError("inputs must be rectangular numeric arrays") from error
if hidden.ndim != 3:
raise ValueError("expected [batch, sequence, feature]")
if projection.ndim != 2:
raise ValueError("expected projection [input_feature, output_feature]")
if hidden.shape[-1] != projection.shape[0]:
raise ValueError("feature axis does not match projection input")
if 0 in hidden.shape or 0 in projection.shape:
raise ValueError("this exercise expects nonempty axes")
return hidden @ projection
hidden = [
[[1.0, 0.0], [0.5, 0.5], [0.0, 1.0]],
[[0.0, 1.0], [0.5, 0.5], [1.0, 0.0]],
]
projection = [[1.0, 0.0, 0.5], [0.0, 1.0, 0.5]]
projected = project_prompts(hidden, projection)
prompt_summaries = projected.mean(axis=1)
print("input", np.shape(hidden), "projection", np.shape(projection))
print("projected", projected.shape)
print("prompt_summaries", prompt_summaries.tolist())
np.testing.assert_allclose(prompt_summaries, [[0.5, 0.5, 0.5]] * 2)
bad_inputs = [
("feature", hidden, np.zeros((3, 3))),
("rank", hidden, [0.0, 0.0]),
("ragged", [[[1.0, 0.0], [0.5]]], projection),
("empty", np.empty((2, 0, 2)), projection),
]
for name, bad_hidden, bad_projection in bad_inputs:
try:
project_prompts(bad_hidden, bad_projection)
except ValueError as error:
print(name, str(error))
else:
raise AssertionError(f"{name} input should have failed")input (2, 3, 2) projection (2, 3)
projected (2, 3, 3)
prompt_summaries [[0.5, 0.5, 0.5], [0.5, 0.5, 0.5]]
feature feature axis does not match projection input
rank expected projection [input_feature, output_feature]
ragged inputs must be rectangular numeric arrays
empty this exercise expects nonempty axesThe two summaries are identical even though the token order differs. That's expected: a shared token projection followed by a mean is insensitive to permutations of the tokens. Shape checks can't make this summary understand order. Later sequence models introduce position information and interactions between tokens for that reason.
Exclude padding from the mean
Real batches often add padding to make sequences equally long. Even zero-valued padding changes an ordinary mean: two real vectors [1, 0] and [0, 1] average to [0.5, 0.5], but adding two zero rows changes that to [0.25, 0.25].
A mask marks real positions. Add a singleton feature axis so one validity bit applies to every coordinate of its token. Sum only real vectors and divide by the count of real tokens. Reject an all-padding prompt because its mean has no valid denominator:
import numpy as np
tokens = np.array([[[1.0, 0.0], [0.0, 1.0], [0.0, 0.0], [0.0, 0.0]]])
valid = np.array([[True, True, False, False]])
def masked_mean(tokens, valid):
if tokens.ndim != 3 or valid.shape != tokens.shape[:2]:
raise ValueError("mask must match [batch, sequence]")
if valid.dtype != np.bool_:
raise ValueError("mask must be boolean")
counts = valid.sum(axis=1, keepdims=True)
if np.any(counts == 0):
raise ValueError("each prompt needs at least one real token")
selected = np.where(valid[:, :, None], tokens, 0.0)
return selected.sum(axis=1) / counts
summary = masked_mean(tokens, valid)
print("plain_mean", tokens.mean(axis=1).tolist())
print("masked_mean", summary.tolist())
np.testing.assert_allclose(summary, tokens[:, :2].mean(axis=1))
try:
masked_mean(tokens, np.zeros_like(valid))
except ValueError as error:
print("empty_prompt", str(error))
else:
raise AssertionError("an all-padding prompt must fail")plain_mean [[0.25, 0.25]]
masked_mean [[0.5, 0.5]]
empty_prompt each prompt needs at least one real tokenThe mask fixes padding bias. A sequence model supplies order-sensitive interactions. Those are separate properties to test. This example assumes finite real-token values. Validate the numeric data you keep as well.
Real token vectors won't be handwritten, and learned projections will be much wider. The invariant is small enough to say aloud. Preserve batch and sequence. Match the feature axis, then let the multiply change its width. Which axis is the one that gets summed away?
Recognize the contract inside a neural-network layer
The same multiply shows up in four shapes. The feature axis is the one that disappears. In this table, N counts rows, D is input width, and P is output width. B and T still mean batch and sequence. Read each contract by naming that summed axis:
Left @ right | What contracts | What survives |
|---|---|---|
[D] @ [D] | matching length | a scalar |
[N, D] @ [D] | feature length | [N] |
[N, D] @ [D, P] | D | [N, P] |
[B, T, D] @ [D, P] | D | [B, T, P] |
One orientation detail matters when you read model code. PyTorch's torch.nn.Linear(in_features, out_features) stores weight as [out_features, in_features], the transpose of our W. It computes input @ weight.T + bias and preserves all input axes except the last, which becomes out_features.[6] So Linear(2, 3) can receive [B, T, 2] directly and return [B, T, 3]. You don't need to flatten the batch or tokens first.
The API stores the weight in a different array orientation; the operation remains the same linear map. Before transposing anything, check how the API names its input and output axes.
Test shapes before code
Keep the convention fixed: rows represent requests or token positions, and the final dimension represents features. Predict each result before checking the sketch:
- A latency feature vector has shape
[3], and a weight vector has shape[3]. What doesx @ wreturn? - A training table
Xhas shape[20, 3], andWhas shape[3, 4]. What doesX @ Wmean and what shape does it produce? - Prompt hidden states have shape
[8, 12, 64], and a projection has shape[64, 16]. Which axes remain unchanged? - Scores have shape
[20, 4]. Why can adding a value of shape[20, 1]be a silent bug when you meant a four-feature bias? - Token vectors have shape
[12, 64]. What transpose is needed to create a pairwise score grid? - For
Hwith shape[2, 3, 2], what scalar doesH[1, 2, 0]select, and what shape does a slice that keeps every prompt and token but fixes feature0have?
Solution sketches
- It returns a scalar because matching feature positions are multiplied and summed.
- It applies one four-output transformation to each of 20 rows, producing
[20, 4]. - Batch size
8and sequence length12stay unchanged; feature size changes from64to16, giving[8, 12, 16]. [20, 1]legally broadcasts one different value across all four output columns in each row. The output shape looks right even though the adjustment is attached to requests instead of features.- Use the transpose
[64, 12]; then[12, 64] @ [64, 12]produces[12, 12]. H[1, 2, 0]selects the first feature of the third token in the second prompt, which is1.0in the running tensor. The feature-0 slice keeps prompt and token axes, giving shape[2, 3].
Carry shape contracts forward
Every model operation now reads as a contract:
| Object | Example shape | Question to ask |
|---|---|---|
| Vector | [2] | What does each feature position mean? |
| Matrix of rows | [3, 2] | What does one row describe? |
| Weight matrix | [2, 3] | Which input feature axis must match? |
| Tensor | [2, 3, 2] | What do batch, sequence, and feature axes mean? |
| Indexing | H[b, t, d] | Which prompt, token position, and feature does this value select? |
| Similarity grid | [3, 3] | Which vector pairs produced these dot products? |
Before running a multiply, predict its output. When it fails, name the meeting axes. When it broadcasts legally, ask whether the repeated value still means what you intended.
You can now name a latency row, stack rows into a matrix, and project a token tensor without guessing shapes. The next question is whether the matrix carries genuinely different feature directions or repeats the same mixture. Matrix rank counts independent directions; singular value decomposition (SVD) reveals their strengths. Those are the next lesson's questions.