PyTorch nn.Linear: shapes, weights, and examples

PyTorch nn.Linear(in_features, out_features) maps the final dimension of an input tensor to a new feature size. It learns a weight matrix and, by default, a bias. A batch of shape (32, 64) passed through Linear(64, 10) produces (32, 10); the layer does not need the batch size in its constructor.

The difficult part is often identifying what the final axis represents. A call can succeed while projecting image width instead of channels, or while preserving sequence positions you intended to combine. This guide makes those choices explicit with small inputs whose outputs can be checked by hand.

The six standalone examples were tested with Python 3.12.14 and PyTorch 2.15.0+cpu. The accompanying checks cover arithmetic, shapes, parameters, errors, and gradients. They do not measure GPU performance or model quality. The official Linear reference documents the constructor and shape contract.

Read the weight matrix in the right direction

For row-oriented input vectors, the operation is y = x @ weight.T + bias. Here Linear(3, 2) owns two rows of three weights. Each row creates one output feature from all three input features.

import torch

x = torch.tensor([[1., 2., 3.], [-1., 0., 2.]], dtype=torch.float64)
layer = torch.nn.Linear(3, 2, dtype=torch.float64)
with torch.no_grad():
    layer.weight.copy_(torch.tensor([[1., 2., -1.], [0., 0.5, 3.]]))
    layer.bias.copy_(torch.tensor([0.5, -1.]))
y = layer(x)
print('Weight shape:', tuple(layer.weight.shape))
print('Output:', y.tolist())
print('Matrix formula matches:', torch.equal(y, x @ layer.weight.T + layer.bias))

Weight shape: (2, 3)
Output: [[2.5, 9.0], [-2.5, 5.0]]
Matrix formula matches: True

For the first input [1, 2, 3], the first output is 1×1 + 2×2 + 3×(-1) + 0.5 = 2.5. The second is 1×0 + 2×0.5 + 3×3 - 1 = 9. The same weights applied to [-1, 0, 2] produce [-2.5, 5].

The stored weight shape is (out_features, in_features), so the transpose matters in this row-vector expression. These outputs also show that Linear does not turn scores into probabilities: values can be negative or exceed one. Choose an activation or loss according to the model's task rather than adding softmax automatically.

We set known parameters under no_grad() to make the arithmetic reproducible. That assignment does not freeze them. The module still registers trainable parameters. A separate check compares its output with functional.linear using the same input, weights, and bias.

Keep batch and sequence dimensions separate

The layer transforms the last dimension and preserves the leading ones. The following all-ones weights make each output feature equal to the sum of the three input features:

import torch

layer = torch.nn.Linear(3, 2, dtype=torch.float64)
with torch.no_grad():
    layer.weight.fill_(1.)
    layer.bias.zero_()
vector = torch.tensor([1., 2., 3.], dtype=torch.float64)
batch = torch.arange(12, dtype=torch.float64).reshape(4, 3)
sequence = torch.arange(24, dtype=torch.float64).reshape(2, 4, 3)
print('Vector:', tuple(layer(vector).shape))
print('Batch:', tuple(layer(batch).shape))
print('Sequence:', tuple(layer(sequence).shape))
print('First sequence:', layer(sequence)[0].tolist())

Vector: (2,)
Batch: (4, 2)
Sequence: (2, 4, 2)
First sequence: [[3.0, 3.0], [12.0, 12.0], [21.0, 21.0], [30.0, 30.0]]

The unbatched vector produces [6, 6]. The four-row batch produces four two-feature vectors. For the sequence input (2, 4, 3), the result (2, 4, 2) retains both batch members and all four positions. It does not average positions, flatten the sequence, or mix the two batch members.

At the first position, [0, 1, 2] sums to 3. At the next, [3, 4, 5] sums to 12. Both output features agree because we deliberately gave their weight rows the same values. In a trained model those rows can learn different projections.

Our checks also change only the second batch member and verify that the first output stays unchanged. A noncontiguous input made by swapping the leading axes produces the corresponding reordered result. Those tests confirm the axis behavior without relying solely on matching shapes.

Count parameters and understand bias=False

For this three-input, two-output layer, six weights plus two bias values give eight trainable scalars. Removing the bias leaves six:

import torch

with_bias = torch.nn.Linear(3, 2)
without_bias = torch.nn.Linear(3, 2, bias=False)
print('With bias:', sum(p.numel() for p in with_bias.parameters()))
print('Without bias:', sum(p.numel() for p in without_bias.parameters()))
print('Absent bias:', without_bias.bias)
print('Saved parameter names:', list(with_bias.state_dict()))
print('No-bias parameter names:', list(without_bias.state_dict()))

With bias: 8
Without bias: 6
Absent bias: None
Saved parameter names: ['weight', 'bias']
No-bias parameter names: ['weight']

In general, the count is in_features × out_features, plus out_features when bias is enabled. Increasing the batch size or sequence length does not add parameters; the same layer is reused at each leading position.

Without a bias, a zero input maps to zero. With a bias, a zero input returns the bias vector. We check both cases. The default parameters are initialized automatically, so a newly created layer will not reproduce our first example unless you assign those values or load matching saved parameters.

The state_dict entries show what is stored for this layer. Loading weights also requires a compatible model definition. Our PyTorch save-and-load guide covers checkpoints and resuming training; keep architecture settings such as feature counts alongside the checkpoint.

Choose what to project in an image tensor

Consider images in (N, C, H, W) order with shape (2, 3, 2, 4). Calling Linear(4, 5) projects width, because width is the final dimension. If you want one output vector per image, flatten the non-batch axes first. For a channel projection at every pixel, move channels to the final axis before applying the layer.

import torch

images = torch.arange(48, dtype=torch.float64).reshape(2, 3, 2, 4)
width_layer = torch.nn.Linear(4, 5, dtype=torch.float64)
image_layer = torch.nn.Linear(24, 5, dtype=torch.float64)
channel_layer = torch.nn.Linear(3, 5, dtype=torch.float64)
print('Width projection:', tuple(width_layer(images).shape))
flat = torch.flatten(images, start_dim=1)
print('Flattened input:', tuple(flat.shape))
print('Whole-image projection:', tuple(image_layer(flat).shape))
channels_last = images.permute(0, 2, 3, 1)
projected = channel_layer(channels_last).permute(0, 3, 1, 2)
print('Per-pixel channel projection:', tuple(projected.shape))

Width projection: (2, 3, 2, 5)
Flattened input: (2, 24)
Whole-image projection: (2, 5)
Per-pixel channel projection: (2, 5, 2, 4)

These are three different operations. The width projection keeps channels and height separate. The whole-image projection joins 3×2×4 = 24 inputs into one vector per image. The channel projection combines the three channel values independently at each of the eight spatial positions, then restores channels-first order.

Our per-pixel test calculates one projected pixel using scalar multiplication and addition, then compares it with the tensor result. It also confirms that flattening preserves the first image's 24 values in order.

The flatten reference explains start_dim. Using start_dim=1 here preserves the batch axis. Flattening from dimension zero would merge images into the same vector. A reshape alone is also not a substitute for moving the channel axis deliberately.

Diagnose common Linear errors

Check the input's shape, dtype, and device alongside the layer's parameters. The following cases deliberately violate the expected final dimension, floating-point type, or minimum tensor rank:

import torch

layer = torch.nn.Linear(3, 2, dtype=torch.float64)
cases = {
    'Wrong final dimension': torch.zeros(2, 4, dtype=torch.float64),
    'Wrong dtype': torch.zeros(2, 3, dtype=torch.float32),
    'Scalar input': torch.tensor(1., dtype=torch.float64),
}
for name, value in cases.items():
    try:
        layer(value)
    except RuntimeError:
        print(name + ': RuntimeError')

Wrong final dimension: RuntimeError
Wrong dtype: RuntimeError
Scalar input: RuntimeError

The first input ends in four values, but the layer expects three. The second uses float32 against float64 parameters in an ordinary call without autocast. The third has no feature axis. All three raise RuntimeError in the tested version. The example prints the exception class because detailed error wording can change between releases.

A matrix-multiplication shape error is a reason to inspect axis meanings, not to change in_features until the program runs. Derive the feature count from the preceding operation. For image inputs, decide between the three projections above before choosing a layer size.

For the dtype case, convert inputs to the intended floating-point dtype or construct the model consistently. In GPU code, inputs and parameters also need compatible devices. These examples run on CPU and do not test CUDA or mixed-precision behavior.

Check gradients separately from evaluation mode

The final example sums both output features and backpropagates. Its fixed weights let us predict each derivative:

import torch

layer = torch.nn.Linear(3, 2, dtype=torch.float64)
with torch.no_grad():
    layer.weight.copy_(torch.tensor([[1., 2., -1.], [0., 0.5, 3.]]))
    layer.bias.copy_(torch.tensor([0.5, -1.]))
x = torch.tensor([1., 2., 3.], dtype=torch.float64, requires_grad=True)
train_output = layer(x)
layer.eval()
y = layer(x)
y.sum().backward()
print('Train/eval outputs equal:', torch.equal(train_output, y))
print('Input gradient:', x.grad.tolist())
print('Weight gradient:', layer.weight.grad.tolist())
print('Bias gradient:', layer.bias.grad.tolist())
with torch.no_grad():
    inference_output = layer(x)
print('no_grad output requires_grad:', inference_output.requires_grad)

Train/eval outputs equal: True
Input gradient: [1.0, 2.5, 2.0]
Weight gradient: [[1.0, 2.0, 3.0], [1.0, 2.0, 3.0]]
Bias gradient: [1.0, 1.0]
no_grad output requires_grad: False

The input gradient [1, 2.5, 2] is the sum of the two weight rows. Each weight row receives [1, 2, 3], the input vector, and each bias receives a gradient of one. The verification code also compares the input gradient with a central finite difference computed from independent Python scalar arithmetic.

The isolated Linear layer produces equal train and eval outputs with unchanged parameters. Calling eval() does not turn off gradient tracking, as the successful backward pass demonstrates. The later no_grad() call produces an output that does not require gradients. Other modules in a complete model can behave differently between training and evaluation; see our guide to eval, no_grad, and inference_mode.

Before using a Linear layer in a larger model, label the input axes, verify the expected output shape, and check one small numerical example. That makes accidental projections easier to catch than a training run that completes with the wrong interpretation of the data.