Broadcasting in PyTorch - ML Experiment: Train & Evaluate

Experiment - Broadcasting

Problem:You want to add a 1D tensor to each row of a 2D tensor using PyTorch. Currently, you try to add them directly but get a size mismatch error.

Current Metrics:Error: RuntimeError: The size of tensor a (3) must match the size of tensor b (4) at non-singleton dimension 1

Issue:The tensors have incompatible shapes for addition. You need to use broadcasting to add the smaller tensor to each row of the larger tensor.

Your Task

Use PyTorch broadcasting to add a 1D tensor of shape (3,) to each row of a 2D tensor of shape (4,3) without errors.

Do not manually expand or repeat the smaller tensor using loops.

Use PyTorch's automatic broadcasting features.

Hint 1

Hint 2

Hint 3

Solution

PyTorch

import torch

# Create a 2D tensor of shape (4, 3)
matrix = torch.tensor([[1, 2, 3],
                       [4, 5, 6],
                       [7, 8, 9],
                       [10, 11, 12]], dtype=torch.float32)

# Create a 1D tensor of shape (3,)
vector = torch.tensor([10, 20, 30], dtype=torch.float32)

# Add vector to each row of matrix using broadcasting
result = matrix + vector

print("Matrix shape:", matrix.shape)
print("Vector shape:", vector.shape)
print("Result shape:", result.shape)
print("Result tensor:")
print(result)

Used PyTorch's automatic broadcasting by adding a 1D tensor directly to a 2D tensor.

Ensured the smaller tensor's shape matches the last dimension of the larger tensor.

Avoided manual expansion or loops.

Results Interpretation

Before: Addition caused a size mismatch error because shapes (4,3) and (3,) were not aligned properly for addition without broadcasting.

After: PyTorch automatically broadcasted the (3,) tensor across the rows of the (4,3) tensor, resulting in successful element-wise addition with shape (4,3).

Broadcasting lets you perform operations on tensors of different shapes by automatically expanding dimensions of size 1, saving you from writing loops or manual expansions.

Bonus Experiment

Try subtracting a 2D tensor of shape (4,1) from a 2D tensor of shape (4,3) using broadcasting.

💡 Hint

PyTorch will broadcast the (4,1) tensor along the last dimension to match (4,3). Check the shapes before subtraction.