Tensor operations and broadcasting

The idea

A neural network is mostly three operations repeated many times: matrix multiplication, element-wise functions and reductions (sums, means). You already know all three from NumPy. This module maps them to PyTorch and points out the handful of differences that trip people up.

Element-wise maths

Tensor maths looks just like NumPy. Prices times units gives revenue per product, and multiplying by 1.11 adds VAT to every price at once:

import torch

prices = torch.tensor([185.0, 320.0, 95.0])
units = torch.tensor([9.0, 10.0, 12.0])
print(prices * units)
print(prices * 1.11)
print(torch.sqrt(units), units.exp().round())

tensor([1665., 3200., 1140.])
tensor([205.3500, 355.2000, 105.4500])
tensor([3.0000, 3.1623, 3.4641]) tensor([  8103.,  22026., 162755.])

Operators work element by element, and most NumPy functions have a PyTorch twin (torch.sqrt, torch.exp, torch.log, torch.abs, torch.clamp), usually also available as methods (units.exp()).

Matrix multiplication

The operation neural networks are built on is matrix multiplication. You've seen it with shop data: each row of order quantities times the price list gives that order's total. It works the same way in PyTorch:

order_qty = torch.tensor([[2.0, 1.0, 3.0],
                          [0.0, 2.0, 0.0]])
price_vec = torch.tensor([95.0, 320.0, 55.0])
print(order_qty @ price_vec)
print(torch.matmul(order_qty, price_vec))

tensor([675., 640.])
tensor([675., 640.])

@ is the same operator as in the NumPy linear-algebra module, with the same shape rule: (n, k) @ (k, m) → (n, m). A neural network layer is exactly X @ W.T + b, as you'll see in the nn.Module lesson.

Reductions along a dimension

PyTorch calls axes dimensions (dim instead of axis), but the rule is the same: the dimension you name disappears.

units = torch.tensor([[120.0, 135.0, 150.0],
                      [ 80.0,  85.0,  90.0]])      # stores × months
print(units.sum())
print(units.sum(dim=0))         # per month
print(units.mean(dim=1))        # per store
print(units.max(dim=1))

tensor(660.)
tensor([200., 220., 240.])
tensor([135.,  85.])
torch.return_types.max(
values=tensor([150.,  90.]),
indices=tensor([2, 2]))

Note max(dim=...): it returns both the maximum values and their positions (indices), like NumPy's max and argmax in one call. Use .values and .indices, or unpack: best, where = units.max(dim=1).

keepdim=True keeps the reduced dimension with size 1, ready for broadcasting:

share = units / units.sum(dim=1, keepdim=True)
print(share.round(decimals=3))

tensor([[0.2960, 0.3330, 0.3700],
        [0.3140, 0.3330, 0.3530]])

Broadcasting

Identical rules to NumPy: compare shapes from the right; sizes must match or be 1.

targets = torch.tensor([[130.0], [85.0]])      # one target per store: shape (2, 1)
print(units - targets)
monthly_price = torch.tensor([95.0, 100.0, 105.0])  # one price per month: shape (3,)
print(units * monthly_price)

tensor([[-10.,   5.,  20.],
        [ -5.,   0.,   5.]])
tensor([[11400., 13500., 15750.],
        [ 7600.,  8500.,  9450.]])

In-place operations

Methods ending in an underscore modify the tensor in place:

w = torch.tensor([1.0, 2.0, 3.0])
w.add_(10)
print(w)
w.mul_(0.5)
print(w)

tensor([11., 12., 13.])
tensor([5.5000, 6.0000, 6.5000])

w.add_(10) changes w; w.add(10) (or w + 10) returns a new tensor. Optimizers update weights in place this way. In your own code, prefer the non-underscore versions unless you have a reason: they're safer (no surprise changes to shared tensors) and they play well with autograd.

Comparisons and masks

Picking values by a condition, like the busy days with more than 60 orders, also works exactly as it does in NumPy:

daily = torch.tensor([42, 38, 51, 47, 60, 75, 68])
busy = daily > 60
print(busy)
print(daily[busy], busy.sum().item(), busy.float().mean().item())
print(torch.where(busy, daily, torch.zeros_like(daily)))

tensor([False, False, False, False, False,  True,  True])
tensor([75, 68]) 2 0.2857142984867096
tensor([ 0,  0,  0,  0,  0, 75, 68])

Boolean masks, torch.where and .sum() of a mask all behave like their NumPy versions. One difference: a boolean tensor must be converted (.float()) before taking its mean.

Why this matters for clean code

Writing model code as whole-tensor operations, rather than Python loops over examples, is what makes it fast on GPUs, and it keeps the code close to the maths on paper: logits = x @ W.T + b instead of three nested loops. The vectorisation habits from the NumPy track transfer directly.

Recap

• Element-wise operators and functions (torch.exp, clamp...) work as in NumPy.
• @ / torch.matmul for matrix products: (n, k) @ (k, m) → (n, m).
• Reductions take dim= (NumPy's axis); keepdim=True keeps the dimension; max(dim) returns values and indices.
• Broadcasting rules are identical to NumPy.
• add_, mul_... modify in place; prefer the non-underscore versions in your own code.
• Masks and torch.where work like NumPy; convert booleans with .float() before mean().