Matrix multiplication: track shapes before calculating

For a matrix product, check the inner dimensions before calculating: an (m, n) matrix multiplied by an (n, p) matrix produces an (m, p) result. In a modelling example, those dimensions often represent samples, input features and outputs.

Shape checks catch many errors early, but a compatible shape does not guarantee correct feature order or business meaning. Track what each axis represents as well as its length.

Define the axes

Let X contain three samples and two features: rows (1, 2), (3, 4) and (5, 6). Let W map those two features to two outputs, with rows (2, -1) and (0, 3).

X has shape (3, 2) and W has shape (2, 2). The product XW therefore has shape (3, 2): one row per sample and one column per output.

For the first sample, the first output is 1 × 2 + 2 × 0 = 2. The second is 1 × (-1) + 2 × 3 = 5. Each output is a dot product between a sample row and a weight column.

Check the product and bias

import numpy as np

x = np.array([[1,2],[3,4],[5,6]])
w = np.array([[2,-1],[0,3]])
bias = np.array([1,-2])
assert x.shape==(3,2) and w.shape==(2,2) and bias.shape==(2,)
product = x@w
assert product.shape==(3,2)
assert np.array_equal(product,[[2,5],[6,9],[10,13]])
output = product+bias
assert np.array_equal(output,[[3,3],[7,7],[11,11]])
try:
    w@x
except ValueError:
    pass
else:
    raise AssertionError('Reversed product should have incompatible inner dimensions')
print({'X_shape':x.shape,'W_shape':w.shape,'product':product.tolist(),
       'output_with_bias':output.tolist()})

The bias has one value per output and broadcasts across the three sample rows. NumPy's matrix multiplication documentation describes the operation behind @. This example uses the locally tested NumPy 2.4.4 runtime.

The mathematics lab records the expected product and a reversed-order failure check.

Distinguish matrix and elementwise multiplication

@ combines the feature axis through sums of products. * multiplies corresponding elements under broadcasting rules. They are different operations even when both happen to return an array.

In a square example, using * instead of @ may not raise an error, which makes an expected-value check useful. Predict one output cell by hand and compare it with the implementation.

Do not fix a shape error by adding a transpose until the code runs. Explain why the transposed axis now represents the intended samples, features or outputs. A syntactically valid transpose can still produce the wrong calculation.

Watch one-dimensional arrays

A vector with shape (2,) is not the same object shape as a row matrix (1, 2) or column matrix (2, 1). NumPy applies specific rules when one-dimensional arrays participate in matrix multiplication.

When a later operation requires an explicit batch dimension, preserve it deliberately. Print or assert shapes at the boundary rather than relying on an example with one sample to behave like a larger batch.

The same caution applies to bias arrays. A shape that broadcasts successfully may add values along an unintended axis if the dimensions happen to match.

Keep a shape ledger

Record each array's name, shape and axis meaning. For this example: X is samples by features, W is features by outputs, bias is outputs and the result is samples by outputs.

For a larger model, extend the ledger through preprocessing, embeddings, batches and predictions. Include feature order and units when relevant. Shape assertions are particularly useful where data enters or leaves a component.

Verify meaning after shape

Two feature columns can be swapped without changing X's shape. The resulting prediction may be wrong even though every dimension check passes. Preserve the fitted feature schema and test a known input-output case.

Exercise: add a fourth sample and then a third output. Predict the required shapes of X, W and bias before editing the arrays. Explain which dimension changes for each modification.

NeuraPath's Data Science course connects linear algebra with model implementation. Tracking shapes and axis meaning makes matrix operations easier to debug and prevents a valid-looking array from concealing the wrong computation.

Continue learning

This article is part of the Mathematics and statistical foundations sequence. Use the neighbouring tasks when you need the prerequisite or the next application.

• Review the prerequisite or neighbouring task in Vectors and dot products through a similarity example.
• Continue with Eigenvectors and PCA through a two-feature dataset.