PyTorch Tensor Shapes: Broadcasting, Reshape, View, Permute and the Errors That Waste Your Time

Page content

PyTorch: Zero to Hero — Step 01

Most PyTorch bugs are not really “AI bugs”.

They are shape bugs.

You expected:

[batch, features]

but actually had:

[batch, sequence, features]

You expected two tensors to line up.

They broadcast instead.

You called view() after permute() and got a contiguity error.

You removed a dimension with squeeze() and accidentally removed the batch dimension too.

Or you reached the familiar message:

RuntimeError: The size of tensor a (...) must match the size of tensor b (...)

This article is about becoming dangerous enough with tensors that these errors stop being mysterious.

We are going to use code heavily.

If you are a programmer coming to PyTorch, the most useful mental model is this:

A tensor is typed multidimensional data plus a shape, a dtype, a device and a memory layout.

The numbers matter.

But when writing PyTorch, the metadata around those numbers often determines whether your program works.


1. Start by interrogating every tensor

Create a tensor:

import torch

x = torch.randn(32, 128)

Before doing anything else, learn to ask it questions:

print(x.shape)
print(x.ndim)
print(x.dtype)
print(x.device)
print(x.numel())
print(x.stride())
print(x.is_contiguous())

Typical output:

torch.Size([32, 128])
2
torch.float32
cpu
4096
(128, 1)
True

For debugging, I often want a tiny helper:

def inspect_tensor(name: str, x: torch.Tensor) -> None:
    print(
        f"{name}: "
        f"shape={tuple(x.shape)}, "
        f"dtype={x.dtype}, "
        f"device={x.device}, "
        f"stride={x.stride()}, "
        f"contiguous={x.is_contiguous()}"
    )

Use it everywhere while developing:

x = torch.randn(8, 16, 64)
inspect_tensor("x", x)
x: shape=(8, 16, 64), dtype=torch.float32, device=cpu, stride=(1024, 64, 1), contiguous=True

That one line often tells you more than staring at the model definition.


2. Shape is part of the type

Python’s type system sees both of these as tensors:

x = torch.randn(32, 128)
y = torch.randn(32, 10, 128)

But semantically they are very different.

I recommend reading shapes as named dimensions.

x: [B, F]
y: [B, T, F]

where:

B = batch
T = time / sequence
F = features

For images you will commonly see:

[B, C, H, W]

For token embeddings:

[B, T, D]

For attention scores:

[B, H, T, T]

where the second H means attention heads, not image height.

The letters are not PyTorch syntax.

They are a way of thinking.

When you read:

x = torch.randn(32, 100, 768)

you should train yourself to ask:

32 what?
100 what?
768 what?

If the answer is:

32 batches
100 tokens
768 embedding dimensions

then mentally write:

[B, T, D]

The code becomes much easier to reason about.


3. Indexing changes shapes

Start with:

x = torch.randn(4, 3, 32, 32)

Think:

[B, C, H, W]

Now inspect several indexing operations:

print(x.shape)
print(x[0].shape)
print(x[:, 0].shape)
print(x[:, :, 0].shape)
print(x[:, :, 0, 0].shape)

Output:

torch.Size([4, 3, 32, 32])
torch.Size([3, 32, 32])
torch.Size([4, 32, 32])
torch.Size([4, 3, 32])
torch.Size([4, 3])

Integer indexing removes a dimension.

Slicing generally preserves it:

print(x[0].shape)     # integer -> removes dim
print(x[0:1].shape)   # slice -> keeps dim
torch.Size([3, 32, 32])
torch.Size([1, 3, 32, 32])

That difference matters enormously.

Suppose a function expects batches:

def process_batch(batch):
    # expected [B, C, H, W]
    print(batch.shape)

This:

process_batch(x[0])

passes:

[C, H, W]

This:

process_batch(x[0:1])

passes:

[1, C, H, W]

If you are debugging code that works with batch size 8 but fails with a single example, check whether you accidentally removed the batch dimension.


4. unsqueeze() is not cosmetic

You will see unsqueeze() constantly in PyTorch code.

x = torch.tensor([1.0, 2.0, 3.0])

print(x.shape)
print(x.unsqueeze(0).shape)
print(x.unsqueeze(1).shape)
torch.Size([3])
torch.Size([1, 3])
torch.Size([3, 1])

Those two shapes represent different things:

[1, 3] = one row containing three values
[3, 1] = three rows containing one value

This becomes important during broadcasting and matrix multiplication.

Example:

x = torch.tensor([1.0, 2.0, 3.0])
row = x.unsqueeze(0)     # [1, 3]
column = x.unsqueeze(1)  # [3, 1]

print(row @ column)   # [1, 3] @ [3, 1] -> [1, 1]
tensor([[14.]])

Reverse it:

print(column @ row)   # [3, 1] @ [1, 3] -> [3, 3]
tensor([[1., 2., 3.],
        [2., 4., 6.],
        [3., 6., 9.]])

Same values.

Very different computation.

Shape is semantics.


5. squeeze() can quietly break your code

Consider:

x = torch.randn(1, 10)

This is:

[B, classes]

with batch size 1.

Now:

print(x.squeeze().shape)
torch.Size([10])

The batch dimension disappeared.

Sometimes that is what you want.

Sometimes it creates a bug that only appears when batch_size == 1.

Prefer specifying the dimension when you know what you mean:

x = torch.randn(1, 10, 1)

print(x.squeeze(-1).shape)   # only remove the trailing size-1 dim
torch.Size([1, 10])

This is much safer than:

x.squeeze()

when dimensions may legitimately have size 1.


6. Broadcasting: useful, powerful, dangerous

PyTorch supports broadcasting.

This lets tensors with different shapes participate in the same operation when their dimensions are compatible.

Start simple:

x = torch.tensor([
    [1.0, 2.0, 3.0],
    [4.0, 5.0, 6.0],
])

bias = torch.tensor([10.0, 20.0, 30.0])

print(x + bias)

Output:

tensor([[11., 22., 33.],
        [14., 25., 36.]])

Shapes:

x    = [2, 3]
bias =    [3]

Line them up from the right:

[2, 3]
   [3]

The final dimensions match.

PyTorch behaves as if the bias existed for both rows.

Conceptually:

[10, 20, 30]
[10, 20, 30]

but it does not need to physically create that repeated tensor.


7. Learn the broadcasting rule

Compare dimensions from right to left.

Two dimensions are compatible if:

they are equal
OR
one of them is 1
OR
one dimension does not exist

We can visualise the rule as a decision tree:

    flowchart TD
    A[Align shapes from the right] --> B{Dimensions match?}
    B -- Yes --> C[Compatible, move left]
    B -- No --> D{Is one dimension 1?}
    D -- Yes --> E[Broadcast the 1 to match, move left]
    D -- No --> F{Does one dimension not exist?}
    F -- Yes --> G[Broadcast to match, move left]
    F -- No --> H[Error: incompatible shapes]
    C --> I{More dimensions?}
    E --> I
    G --> I
    I -- Yes --> B
    I -- No --> J[Success]
  

Example:

x: [32, 10, 128]
y:         [128]

Compatible.

Example:

x: [32, 10, 128]
y:      [10, 128]

Compatible.

Example:

x: [32, 10, 128]
y:  [1,  1, 128]

Compatible.

But:

x: [32, 10, 128]
y:     [32, 128]

is not compatible in the way people often expect.

Align from the right:

[32, 10, 128]
    [32, 128]

Compare:

128 == 128  ✓
10  != 32   ✗

That produces a shape error.

The cure may be to insert the missing dimension:

y = torch.randn(32, 128)
y = y.unsqueeze(1)

print(y.shape)
torch.Size([32, 1, 128])

Now:

x: [32, 10, 128]
y: [32,  1, 128]

broadcasts over the sequence dimension.


8. A broadcasting bug that does not crash

The worst shape bugs are not the ones that throw exceptions.

They are the ones that run.

Consider:

predictions = torch.randn(32, 1)
targets = torch.randn(32)

loss = (predictions - targets) ** 2

print(loss.shape)

What shape do you expect?

Many programmers expect:

[32, 1]

Actual result:

[32, 32]

Why?

predictions: [32, 1]
targets:        [32]

Align from the right:

[32,  1]
     [32]

The 1 broadcasts to 32, and the missing leading dimension broadcasts too.

You have accidentally compared every prediction against every target.

Your code runs.

Your loss is nonsense.

Fix it explicitly:

targets = targets.unsqueeze(1)

print(predictions.shape)   # [32, 1]
print(targets.shape)       # [32, 1]

Then:

loss = (predictions - targets) ** 2
print(loss.shape)          # [32, 1]

This is why shape assertions are worth writing.


9. Assert your assumptions

Programmers write assertions for invariants everywhere else.

Do the same in ML code.

def mse(predictions: torch.Tensor, targets: torch.Tensor) -> torch.Tensor:
    assert predictions.shape == targets.shape, (
        f"shape mismatch: predictions={predictions.shape}, "
        f"targets={targets.shape}"
    )

    return ((predictions - targets) ** 2).mean()

Now the previous bug fails immediately:

predictions = torch.randn(32, 1)
targets = torch.randn(32)

mse(predictions, targets)

Instead of silently producing bad training data, you get an error at the boundary where the assumption was violated.

For larger models, I like assertions such as:

assert x.ndim == 3
assert x.shape[-1] == embedding_dim
assert logits.shape[:-1] == targets.shape
assert attention_mask.dtype == torch.bool

Treat tensor shape as part of the function contract.


10. reshape() changes interpretation, not data

Create 24 sequential values:

x = torch.arange(24)
print(x)
print(x.shape)
tensor([ 0,  1,  2, ..., 23])
torch.Size([24])

Reshape:

y = x.reshape(4, 6)

print(y)
print(y.shape)
torch.Size([4, 6])

Or:

z = x.reshape(2, 3, 4)
print(z.shape)
torch.Size([2, 3, 4])

The number of elements must remain compatible:

print(x.numel())
24

So this works:

x.reshape(3, 8)

This does not:

x.reshape(5, 5)

because 25 slots cannot contain 24 elements.


11. Let PyTorch infer one dimension with -1

Instead of calculating every dimension manually:

x = torch.randn(32, 3, 28, 28)

Flatten each image:

flat = x.reshape(32, -1)

print(flat.shape)
torch.Size([32, 2352])

because:

3 × 28 × 28 = 2352

A more robust pattern is:

flat = x.reshape(x.shape[0], -1)

Now the batch size does not need to be hard-coded.

You will use this constantly.


12. flatten() often communicates intent better

These are similar:

flat = x.reshape(x.shape[0], -1)

and:

flat = torch.flatten(x, start_dim=1)

For:

x = torch.randn(32, 3, 28, 28)

both produce:

[32, 2352]

I often prefer:

torch.flatten(x, 1)

because it says what I mean:

preserve the batch dimension and flatten everything after it.


13. view() and reshape() are not identical

This catches many PyTorch programmers eventually.

For a simple contiguous tensor:

x = torch.arange(24)

print(x.view(4, 6).shape)
print(x.reshape(4, 6).shape)

Both work.

But tensor operations can change how the data is laid out logically in memory.

Consider:

x = torch.arange(12).reshape(3, 4)
y = x.transpose(0, 1)

print(x.shape)
print(y.shape)
print(x.is_contiguous())
print(y.is_contiguous())

Typical output:

torch.Size([3, 4])
torch.Size([4, 3])
True
False

Now try:

y.view(-1)

You may get:

RuntimeError: view size is not compatible with input tensor's size and stride ...

Why?

view() requires a memory layout compatible with the requested view.

reshape() is more flexible:

flat = y.reshape(-1)
print(flat)

It returns the requested shape, copying when necessary.

A useful practical rule:

Use view() when you deliberately care about view semantics.
Use reshape() when you mainly care about the resulting shape.

But understand that reshape() may return a view or may copy.

Do not write logic that depends on which one happened.


14. What does contiguous mean?

Start with:

x = torch.arange(12).reshape(3, 4)

Its logical rows correspond naturally to its underlying storage.

print(x.stride())
(4, 1)

Moving one row means skipping four elements.

Moving one column means moving one element.

Transpose it:

y = x.transpose(0, 1)

print(y.shape)
print(y.stride())

Now you may see:

(1, 4)

The data itself did not need to be physically rearranged.

PyTorch changed how the dimensions map onto the existing storage.

That is efficient.

It also explains why some view operations cannot be performed directly afterward.

If you explicitly need contiguous storage:

y = y.contiguous()
print(y.is_contiguous())
True

Then:

y.view(-1)

works.

Do not sprinkle .contiguous() everywhere merely to silence errors.

Understand why the tensor stopped being contiguous first.

Copies cost memory and time.


15. permute() changes dimension order

This is essential for image models, sequence models and attention code.

Suppose data arrives in image-last format:

[B, H, W, C]

Example:

images = torch.randn(32, 224, 224, 3)

Many PyTorch convolution operations expect:

[B, C, H, W]

Use:

images = images.permute(0, 3, 1, 2)

print(images.shape)
torch.Size([32, 3, 224, 224])

The arguments mean:

old dimensions: 0 1 2 3
old meaning:    B H W C
new order:      0 3 1 2
new meaning:    B C H W

For transformer code you may see:

x = torch.randn(32, 128, 768)
[B, T, D]

and then:

x = x.reshape(32, 128, 12, 64)
[B, T, H, Dh]

followed by:

x = x.permute(0, 2, 1, 3)
[B, H, T, Dh]

That pattern will matter when we build attention later in the series.

We can diagram the transformation:

    flowchart LR
    subgraph reshape
        A["[B, T, D]"] --> A1["[B, T, H, Dh]"]
    end
    subgraph permute
        A1 --> A2["[B, H, T, Dh]"]
    end
    style A fill:#f9f,stroke:#333
    style A1 fill:#bbf,stroke:#333
    style A2 fill:#bfb,stroke:#333
  

16. transpose() versus permute()

transpose() swaps two dimensions:

x = torch.randn(2, 3, 4)

y = x.transpose(1, 2)

print(y.shape)
torch.Size([2, 4, 3])

permute() specifies the complete order:

y = x.permute(0, 2, 1)

Same result here:

[2, 4, 3]

For two-dimensional matrices:

x.T

is convenient.

For model code with several dimensions, explicit dimension names in comments help:

# [B, T, H, Dh] -> [B, H, T, Dh]
x = x.permute(0, 2, 1, 3)

That comment is worth more than it looks.


17. Matrix multiplication has shape rules too

Consider:

x = torch.randn(32, 128)
w = torch.randn(128, 64)

out = x @ w
print(out.shape)
torch.Size([32, 64])

Think:

[B, Din] @ [Din, Dout] -> [B, Dout]

The inner dimensions must match.

This fails:

x = torch.randn(32, 128)
w = torch.randn(256, 64)

x @ w

because:

128 != 256

When you see:

RuntimeError: mat1 and mat2 shapes cannot be multiplied

print the two shapes immediately.

print("x:", x.shape)
print("w:", w.shape)

Do not start changing layer sizes at random.

Trace where the unexpected dimension entered the pipeline.


18. Batched matrix multiplication

PyTorch’s @ / matmul handles batches too.

x = torch.randn(32, 10, 128)
w = torch.randn(128, 64)

out = x @ w

print(out.shape)
torch.Size([32, 10, 64])

Think:

[B, T, Din] @ [Din, Dout]
[B, T, Dout]

PyTorch applies the matrix multiplication across the leading dimensions.

This is the basis of many neural-network layers.


19. Attention is mostly shape manipulation plus matrix multiplication

Here is a preview of where this is going.

Suppose:

batch_size = 2
sequence_length = 8
embedding_dim = 64
num_heads = 4
head_dim = embedding_dim // num_heads

q = torch.randn(batch_size, sequence_length, embedding_dim)
k = torch.randn(batch_size, sequence_length, embedding_dim)
v = torch.randn(batch_size, sequence_length, embedding_dim)

Split embeddings into heads:

q = q.reshape(batch_size, sequence_length, num_heads, head_dim)
k = k.reshape(batch_size, sequence_length, num_heads, head_dim)
v = v.reshape(batch_size, sequence_length, num_heads, head_dim)

Shapes:

[B, T, H, Dh]

Move the head dimension:

q = q.permute(0, 2, 1, 3)
k = k.permute(0, 2, 1, 3)
v = v.permute(0, 2, 1, 3)

Now:

[B, H, T, Dh]

Compute attention scores:

scores = q @ k.transpose(-2, -1)

print(scores.shape)
[B, H, T, T]

Already, without knowing much about transformers, you can derive why the shape is [B, H, T, T].

That is the skill we are building.


20. Dtypes create another class of bugs

Tensor shape is only part of the contract.

Inspect dtype too:

x = torch.tensor([1, 2, 3])
y = torch.tensor([1.0, 2.0, 3.0])

print(x.dtype)
print(y.dtype)

Typically:

torch.int64
torch.float32

Many model operations expect floating-point tensors.

Class labels for CrossEntropyLoss, however, are normally integer class indices.

Example:

logits = torch.randn(4, 10)
targets = torch.tensor([1, 5, 3, 0], dtype=torch.long)

This combination is intentional:

logits  -> float
labels  -> integer

Blindly converting every tensor to float can therefore be just as wrong as leaving everything integer.

Again: dtype is part of the semantic type.


21. Device mismatch is the shape error’s close relative

Later we will use accelerators.

A common failure looks like:

Expected all tensors to be on the same device ...

The simplest prevention is to move related tensors together.

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

x = torch.randn(32, 128).to(device)
w = torch.randn(128, 64).to(device)

out = x @ w

Debug with:

print(x.device)
print(w.device)

For models:

model = model.to(device)
inputs = inputs.to(device)
targets = targets.to(device)

A tensor’s device belongs in the same mental checklist as its shape and dtype.


22. A reusable tensor debugger

For real projects, make inspection cheap.

from dataclasses import dataclass

import torch


@dataclass(frozen=True)
class TensorInfo:
    shape: tuple[int, ...]
    dtype: torch.dtype
    device: torch.device
    stride: tuple[int, ...]
    contiguous: bool
    requires_grad: bool


def tensor_info(x: torch.Tensor) -> TensorInfo:
    return TensorInfo(
        shape=tuple(x.shape),
        dtype=x.dtype,
        device=x.device,
        stride=x.stride(),
        contiguous=x.is_contiguous(),
        requires_grad=x.requires_grad,
    )

Use it:

x = torch.randn(4, 8, 16, requires_grad=True)
print(tensor_info(x))

Or create a contract helper:

def require_shape(x: torch.Tensor, *shape: int | None) -> None:
    if x.ndim != len(shape):
        raise ValueError(
            f"expected {len(shape)} dimensions, got {x.ndim}: {tuple(x.shape)}"
        )

    for dim, (actual, expected) in enumerate(zip(x.shape, shape)):
        if expected is not None and actual != expected:
            raise ValueError(
                f"dimension {dim}: expected {expected}, got {actual}; "
                f"full shape={tuple(x.shape)}"
            )

Now:

x = torch.randn(32, 100, 768)

require_shape(x, None, None, 768)

passes regardless of batch and sequence length.

This:

bad = torch.randn(32, 100, 512)
require_shape(bad, None, None, 768)

fails close to the source of the problem.


23. Debug shape errors systematically

When a PyTorch operation fails because of dimensions, do not guess.

Use a process.

    flowchart TD
    A[Error!] --> B[Print shapes of all operands]
    B --> C["Attach semantic names (B, T, D)"]
    C --> D[Derive expected operation shape manually]
    D --> E{Shape mismatch?}
    E -- Yes --> F[Find where the wrong dimension appeared upstream]
    F --> G[Fix interpretation, not just a reshape]
    G --> H[Assert the corrected invariant]
    E -- No --> I[Check dtype, device, contiguity]
    I --> J[Problem solved?]
  

Step 1: print the operands

print(a.shape)
print(b.shape)

Step 2: attach semantic names

a = [B, T, D]
b = [?, ?]

Step 3: derive the expected operation manually

For matrix multiplication:

[..., M, K] @ [..., K, N] -> [..., M, N]

For broadcasting, align dimensions from the right.

Step 4: find where the bad dimension first appeared

Do not simply patch the failing line with:

reshape(...)

unless you understand why the data should have that interpretation.

A successful reshape can hide a conceptual bug.

Step 5: assert the corrected invariant

Once you know the intended shape:

assert x.shape[-1] == expected_features

Make the bug harder to reintroduce.


24. Debugging example: classifier shape mismatch

Imagine:

batch_size = 16
channels = 3
height = 32
width = 32

images = torch.randn(batch_size, channels, height, width)

You write a linear layer expecting flattened images:

weights = torch.randn(3 * 32 * 32, 10)

Then accidentally do:

logits = images @ weights

This fails because the image tensor is still four-dimensional.

Inspect:

print(images.shape)   # [16, 3, 32, 32]
print(weights.shape)  # [3072, 10]

What operation did we intend?

[B, 3072] @ [3072, 10]

So make that shape explicit:

flat = images.flatten(start_dim=1)

print(flat.shape)     # [16, 3072]

Now:

logits = flat @ weights
print(logits.shape)   # [16, 10]

That is not “fixing a PyTorch error”.

It is correcting the data representation to match the mathematical operation we intended.


25. Debugging example: image channel order

A NumPy/OpenCV-style pipeline may give you:

images = torch.randn(32, 224, 224, 3)

Your convolution expects:

[B, C, H, W]

but you have:

[B, H, W, C]

The values are fine.

The dimension meaning is wrong.

Fix:

images = images.permute(0, 3, 1, 2)

assert images.shape == (32, 3, 224, 224)

This is why debugging tensor code requires understanding both shape and what each axis means.


26. Debugging example: sequence length versus embedding size

Suppose a transformer block expects:

[B, T, D]

with:

embedding_dim = 768

Add an explicit guard:

def transformer_block(x: torch.Tensor) -> torch.Tensor:
    if x.ndim != 3:
        raise ValueError(f"expected [B, T, D], got {tuple(x.shape)}")

    if x.shape[-1] != 768:
        raise ValueError(
            f"expected embedding dimension 768, got {x.shape[-1]}"
        )

    return x

If upstream code accidentally transposes sequence and embedding dimensions:

x = torch.randn(4, 768, 128)

you now fail immediately instead of discovering the mistake twelve operations later.


27. Do not solve every problem with reshape()

This code is suspicious:

try:
    output = model(x)
except RuntimeError:
    x = x.reshape(...)
    output = model(x)

reshape() is not a type cast for tensors.

It changes the interpretation of the same elements.

Before reshaping, you should be able to state:

Current semantic shape: [B, C, H, W]
Required semantic shape: [B, C*H*W]
Reason: the next operation is a fully connected projection over all image features.

If you cannot explain the transformation, do not make it merely because the element counts line up.


28. A useful development pattern: annotate shapes in code

Tensor-heavy code becomes dramatically easier to maintain when important transformations are annotated.

Instead of:

q = q.reshape(b, t, h, d).permute(0, 2, 1, 3)

write:

# q: [B, T, D]
q = q.reshape(batch_size, seq_len, num_heads, head_dim)
# q: [B, T, H, Dh]
q = q.permute(0, 2, 1, 3)
# q: [B, H, T, Dh]

This is especially valuable in:

  • attention
  • CNN feature pipelines
  • recurrent models
  • batching code
  • multimodal models
  • loss computation

The comments are not noise.

They document the hidden type system of the program.


29. Mini challenge: find the bug before running it

What shape does this produce?

x = torch.randn(8, 12, 64)
bias = torch.randn(12, 1)

y = x + bias

Write the dimensions aligned from the right:

x:    [8, 12, 64]
bias:    [12,  1]

This works.

Result:

[8, 12, 64]

Now change bias:

bias = torch.randn(8, 64)

Does this broadcast?

Align:

[8, 12, 64]
   [8, 64]

Compare from the right:

64 == 64
12 != 8

It fails.

If bias really represents one vector per batch item, what shape do we need?

bias = bias.unsqueeze(1)

Now:

[8, 12, 64]
[8,  1, 64]

and broadcasting does what we intended.


30. Build a shape-safe linear projection

Let’s finish with something close to real model code.

import torch


def linear(
    x: torch.Tensor,
    weight: torch.Tensor,
    bias: torch.Tensor | None = None,
) -> torch.Tensor:
    if x.ndim < 2:
        raise ValueError(
            f"x must have at least 2 dimensions [..., features], got {x.shape}"
        )

    if weight.ndim != 2:
        raise ValueError(
            f"weight must be [in_features, out_features], got {weight.shape}"
        )

    in_features, out_features = weight.shape

    if x.shape[-1] != in_features:
        raise ValueError(
            f"input feature mismatch: x has {x.shape[-1]}, "
            f"weight expects {in_features}"
        )

    if bias is not None and bias.shape != (out_features,):
        raise ValueError(
            f"bias must be [{out_features}], got {tuple(bias.shape)}"
        )

    y = x @ weight

    if bias is not None:
        y = y + bias

    return y

Test a normal batch:

x = torch.randn(32, 128)
weight = torch.randn(128, 64)
bias = torch.randn(64)

out = linear(x, weight, bias)

print(out.shape)
torch.Size([32, 64])

Now sequence data:

x = torch.randn(32, 100, 128)

out = linear(x, weight, bias)

print(out.shape)
torch.Size([32, 100, 64])

Because matrix multiplication and bias broadcasting operate over the trailing feature dimensions, the same function works for both.

Now compare that with what nn.Linear will eventually do for us.

The abstraction will make much more sense because we already understand the tensor operation underneath it.


The tensor debugging checklist

When a PyTorch operation fails, inspect these before changing code:

print("shape:", x.shape)
print("ndim:", x.ndim)
print("dtype:", x.dtype)
print("device:", x.device)
print("stride:", x.stride())
print("contiguous:", x.is_contiguous())
print("requires_grad:", x.requires_grad)

Then ask:

What does every dimension mean?
What shape does the next operation expect?
Is broadcasting occurring?
Did indexing remove a dimension?
Did squeeze remove a dimension?
Did permute change the layout?
Am I using view() on a non-contiguous tensor?
Are my operands on the same device?
Are the dtypes appropriate for the operation?

That checklist will solve a surprising fraction of PyTorch problems.


What you should now be able to read

This should no longer look like random tensor punctuation:

B, T, D = 32, 128, 768
H = 12
Dh = D // H

x = torch.randn(B, T, D)

q = x.reshape(B, T, H, Dh)
q = q.permute(0, 2, 1, 3)

assert q.shape == (B, H, T, Dh)

You should be able to derive every transformation.

That ability matters more than memorizing PyTorch methods.


Next: gradients without magic

In Step 0 we used:

loss.backward()

without fully opening the box.

Now that tensors and tensor operations are less mysterious, the next post will tackle one of the most searched and misunderstood parts of PyTorch:

autograd and gradient debugging.

We will cover:

requires_grad
.grad
.grad_fn
leaf tensors
zeroing gradients
in-place operation errors
detach()
no_grad()
NaN gradients
anomaly detection
"Trying to backward through the graph a second time"

And, as with this post, we will learn it by breaking code and fixing it.