Tensors and shapes for AI engineers
A tensor is an array with any number of dimensions: a scalar has 0, a vector 1, a matrix 2, and a batch of matrices 3 or more. Text, images and audio are unstructured data, so representing them takes more dimensions: a batch of sentences is (batch, tokens, embedding size), and a batch of images is (batch, channels, height, width). Underneath, a tensor is one flat block of memory plus a shape and strides, and most errors you'll debug in AI code are shape mismatches.
From arrays to tensors
| Rank | Name | Example | Shape |
|---|---|---|---|
| 0 | Scalar | A temperature value | () |
| 1 | Vector | One sentence embedding | (384,) |
| 2 | Matrix | A batch of embeddings | (32, 384) |
| 3 | 3D tensor | Token embeddings for a batch of sentences | (32, 128, 384) |
| 4 | 4D tensor | A batch of RGB images | (32, 3, 224, 224) |
Structured data, like a spreadsheet, fits in two dimensions: rows and columns. Unstructured data needs more: a sentence is a sequence of tokens, each token is a vector, and you process many sentences at once.
Tensors as memory
Memory is one long line of bytes, so a tensor is stored flat, with a shape and strides that say how to find each element. For a row-major tensor of shape (2, 3, 4), element (i, j, k) sits at flat index i × 12 + j × 4 + k. The strides are (12, 4, 1).
def strides(shape: tuple[int, ...]) -> tuple[int, ...]:
out, step = [], 1
for dim in reversed(shape):
out.append(step)
step *= dim
return tuple(reversed(out))
def flat_index(index: tuple[int, ...], shape: tuple[int, ...]) -> int:
return sum(i * s for i, s in zip(index, strides(shape)))
shape = (2, 3, 4)
print("strides:", strides(shape))
print("flat index of (1, 2, 3):", flat_index((1, 2, 3), shape))
This is why transposing a tensor can be free (only the strides change), and why some operations need a contiguous copy first. It's also why iterating along the last dimension is fast: those elements sit next to each other in memory.
Working with shapes in NumPy
import numpy as np
batch = np.random.rand(32, 128, 384)
sentence_vecs = batch.mean(axis=1)
normed = sentence_vecs / np.linalg.norm(sentence_vecs, axis=1, keepdims=True)
similarity = normed @ normed.T
print(batch.shape, sentence_vecs.shape, similarity.shape)
Step by step: a batch of 32 sentences with 128 tokens and 384 dimensions per token, averaged over tokens (mean pooling) to one vector per sentence, normalised to length 1, then multiplied by its own transpose to get a 32 × 32 matrix of cosine similarities.
Broadcasting
Broadcasting lets NumPy and PyTorch combine tensors of different shapes by stretching dimensions of size 1.
(32, 384) / (32, 1) divides every row by its own norm.
It's convenient, and it's also how silent bugs happen when a dimension you expected to be 32 is 1.
Reading shape errors
RuntimeError: mat1 and mat2 shapes cannot be multiplied (32x384 and 768x10) means the input has 384 features and the layer expects 768.
Usually the embedding model and the layer were built for different models.
The fix starts with printing shapes at each step.
Name your dimensions in variable names (batch_tokens_dim) or with tools that check shapes, because shapes are the type system of AI code.
Check your understanding
Predict the output shape of each line in the NumPy example before running it.
From the floor
In the AI data structures lecture, tensors came right after the bridge from discrete to continuous values, because once a word is a vector of decimal numbers, a sentence is a matrix and a batch of sentences is a 3D tensor, and everything an LLM computes is operations on those.