Broadcasting (PyTorch)

This note applies the general broadcasting rule from Broadcasting to real PyTorch code — the shape-alignment gotchas that trip up most beginners in actual practice.

The Rule, Applied in PyTorch

import torch

a = torch.ones(4, 3)     # a batch of 4 samples, 3 features
b = torch.ones(3)          # one bias per feature

print((a + b).shape)   # (4, 3) -- b is broadcast across all 4 rows, exactly as expected

The Classic Silent Bug: Unintended Shape Mismatch

predictions = torch.tensor([1.0, 2.0, 3.0])        # shape (3,)
targets = torch.tensor([[1.0], [2.0], [3.0]])        # shape (3, 1) -- easy to create by accident!

diff = predictions - targets
print(diff.shape)   # (3, 3) -- NOT (3,)! Broadcasting silently created a 3x3 outer-product-like result

# The fix: make shapes match explicitly before subtracting
targets_flat = targets.squeeze()   # shape (3,)
diff_correct = predictions - targets_flat
print(diff_correct.shape)   # (3,) -- correct

This exact pattern — one tensor accidentally shaped \((n,1)\) instead of \((n,)\), often from a careless .reshape(), a slicing operation, or a loss function's output — is one of the most common silent bugs in real PyTorch training code, since it produces no error, just a wrong, larger-than-expected result.

Debugging Broadcasting Issues

# Always print shapes before an operation you're unsure about
print(f"predictions: {predictions.shape}, targets: {targets.shape}")
result = predictions - targets
print(f"result: {result.shape}")   # if this is unexpectedly large, broadcasting likely misfired

Common Mistakes

• Not checking tensor shapes before an operation, especially after slicing, reduction operations (.sum(), .mean() with keepdim), or loading batched data — these are common sources of an unexpected extra dimension of size 1.
• Using .squeeze() without specifying a dimension when only one specific dimension should be removed — an unqualified .squeeze() removes every size-1 dimension, which can unexpectedly collapse a batch dimension of size 1.

Interview Relevance

Q: "You compute a loss and it's much larger than expected, with no error thrown. What's a common broadcasting-related cause?" One tensor (often the labels or targets) accidentally has an extra dimension of size 1 — e.g. shape \((n,1)\) instead of \((n,)\) — causing an element-wise operation against a shape-\((n,)\) tensor to silently broadcast into an unintended \((n,n)\) result, rather than the expected \((n,)\) element-wise comparison. Checking both tensors' .shape before the operation immediately reveals this.

Practice Question

What shape would result from adding a tensor of shape (5,) to a tensor of shape (5, 1)? Is this likely to be the intended behavior?

Related DL Notes

• Broadcasting
• Tensor Operations
• PyTorch Loss Functions