Tensor operations and reductions

Elementwise, matrix and reduction operations

Elementwise operations such as +, *, torch.exp and torch.relu work position by position, so the inputs must have the same shape or be broadcastable. A frequent mistake is using * when you mean matrix multiplication: * is elementwise. Matrix multiplication is @ or torch.matmul, and the inner sizes must match: (m, n) @ (n, p) gives (m, p).

Reductions such as sum, mean, max and argmax collapse a dimension. The dim argument names the dimension that disappears. For a (batch, features) tensor, x.mean(dim=0) averages over the batch and returns one value per feature; x.mean(dim=1) returns one value per example. Add keepdim=True to keep the collapsed dimension with size 1, which helps later broadcasting.

Methods ending in an underscore, such as add_ or zero_, modify the tensor in place. They save memory but can break autograd if the original values were needed for gradients. Use .item() to read a one-element tensor as a Python number.

Seeing dim in action

Each comment gives the result. Notice that sum(dim=0) returns three numbers (one per column) because dimension 0, the rows, is the one that collapses.

import torch

x = torch.tensor([[1.0, 2.0, 3.0],
                  [4.0, 5.0, 6.0]])   # shape (2, 3)
w = torch.ones(3, 2)                  # shape (3, 2)

print(x * 2)               # elementwise: [[2, 4, 6], [8, 10, 12]]
print(x @ w)               # (2,3) @ (3,2) -> (2,2): [[6, 6], [15, 15]]
print(x.sum(dim=0))        # rows collapse:    [5., 7., 9.]
print(x.sum(dim=1))        # columns collapse: [6., 15.]
print(x.mean(dim=1, keepdim=True).shape)   # torch.Size([2, 1])
print(x.max().item())      # 6.0

The dim trap

Students often read dim=0 as "per row". It is the opposite: dim=0 removes the row dimension, so you get one result per column. When unsure, print .shape before and after the reduction; the missing size tells you which dimension was reduced.

Quick check: `x` has shape `(32, 10)` (32 examples, 10 features). What is the shape of `x.sum(dim=1)`?

• (10,)
• (1, 10)
• (32,)
• (32, 1)