Read a tensor before training a model

Outcome: calculate a small batched linear transformation and explain its shape, dtype and axes. Prerequisite: numerical arrays, functions and basic supervised learning from the earlier courses.

Open the diagram at full size

Read the diagram explanation

A two by three input times three by two transposed weights yields two by two outputs; bias repeats by row.

A neural network begins with ordinary numbers arranged according to a contract. The arrangement tells the program which numbers belong to one observation and which numbers are features of that observation. If that contract is wrong, a program can produce valid arithmetic for the wrong question. Reading shapes is therefore part of understanding a model, before any discussion of accuracy.

Give each axis a meaning

A tensor is a numerical container with a shape. A scalar, such as one loss value, has shape (). A vector with three entries has shape (3,). A matrix with two rows and three columns has shape (2, 3). The number of axes is sometimes called the tensor's rank; this use of rank is different from matrix rank in linear algebra.

For tabular input, we use shape (B, D): B is the number of observations in the current batch and D is the number of features per observation. A batch groups observations for one computation. It does not create a new feature. If there are 12 observations with three measurements each, the shape is (12, 3), even if the training dataset contains thousands of observations.

A tensor does not know that a column means temperature. Keep the feature order alongside the data. Reversing two equally sized columns preserves shape while changing the question asked of every learned weight. Numerical shape checks and semantic checks are both necessary.

Work through one linear map

Let X name this two-observation input:

X = [[ 1,  2, 0],
     [ 0, -1, 3]]

We want two output numbers per observation. Let W be the weights, with one row for each output, and b be its two biases:

W = [[ 2, -1, 0.5],
     [-1,  0, 1.0]]
b = [0.5, -0.5]

A weight multiplies an input feature; a bias adds a value independently of the input. For each observation, multiply corresponding input and weight entries, add the products, then add the output's bias. The first observation's first output is 1×2 + 2×(-1) + 0×0.5 + 0.5 = 0.5. Its second output is 1×(-1) + 2×0 + 0×1 - 0.5 = -1.5.

For the second observation, the first output is 0×2 + (-1)×(-1) + 3×0.5 + 0.5 = 3. The second is 0×(-1) + (-1)×0 + 3×1 - 0.5 = 2.5. The output is therefore [[0.5, -1.5], [3.0, 2.5]], with shape (2, 2).

The compact expression is Z = X @ W.T + b. Z names the output, @ is matrix multiplication, and .T exchanges the two axes of this matrix. Shapes tell the same story: (2, 3) @ (3, 2) produces (2, 2). PyTorch's Linear(3, 2) stores its weights in shape (2, 3), matching the two rows of weights above. Although called a linear layer, its optional bias makes the transformation affine.

Check the arithmetic in PyTorch

import torch

X = torch.tensor([[1, 2, 0], [0, -1, 3]], dtype=torch.float32)
W = torch.tensor([[2, -1, 0.5], [-1, 0, 1]], dtype=torch.float32)
b = torch.tensor([0.5, -0.5], dtype=torch.float32)
Z = X @ W.T + b
expected = torch.tensor([[0.5, -1.5], [3.0, 2.5]])
assert X.shape == (2, 3)
assert Z.shape == (2, 2)
assert torch.allclose(Z, expected)
print(Z)

dtype describes the numerical representation. float32 is our usual training type; float64 provides extra precision for the tiny numerical gradient check later. Whole-number input written without a dtype may become an integer tensor, which is unsuitable for ordinary floating-point layer computations. Class-index labels will deliberately use an integer dtype in Lesson 02. Do not turn every tensor into floating point indiscriminately.

The device says where the tensor lives. This course uses the CPU. Compatible shapes alone are insufficient when operands have incompatible dtypes or devices. Check those properties when reading an error, instead of trying arbitrary reshapes.

Broadcasting can hide a wrong comparison

Adding b, shape (2,), to Z, shape (2, 2), repeats the same two biases across observations. This is intentional broadcasting. Broadcasting aligns dimensions from the right; dimensions must match or one must have size one.

Now imagine predictions of shape (3, 1) and targets of shape (3,). Subtraction can produce shape (3, 3), comparing each prediction with several targets. The arithmetic succeeds, but observation pairing is lost. For one-output regression, keep both tensors in shape (B, 1) and assert equality before computing the loss.

Likewise, reshape rearranges the view of stored values; it does not establish what an axis means. A transpose exchanges axes. Confusing those operations can scramble observations while retaining a plausible element count. Write the intended axis meanings first, then choose the operation.

Practice

Complete these before opening the solutions.

• 01A · A new observation: use X = [[2, 0, -1]] with the worked weights and biases. Predict the shape and calculate both outputs by hand; then verify them with PyTorch.
• 01B · A broadcasting trap: create predictions [[1.0], [3.0], [5.0]] and targets [2.0, 4.0, 6.0]. Inspect the subtraction's shape. Repair the target shape and explain which comparisons now occur.
• 01C · Count the adjustable numbers: a layer maps four features to three outputs with a bias. State the input, weight, bias and output shapes for a batch of five, and count its learned scalar parameters.

Exit check

Without running code, explain why Linear(3, 2) accepts (7, 3) and returns (7, 2). Then make your own one-observation example and verify each output independently. Passing these checks demonstrates shape reasoning for this layer; it does not yet demonstrate that you can train a useful network.

Further reading: PyTorch 2.9 Linear, tensor construction, and broadcasting semantics.