TopolearnNeural networks are compositions of exactly these operations, so reading model code means reading tensor math.
Arithmetic operators work element-wise: a + b, a * b, a ** 2, plus functions like torch.exp, torch.sqrt and torch.relu. Shapes must match or be broadcastable, with the same rules as NumPy: align shapes from the right; each dimension must be equal or 1. A bias of shape (n,) therefore adds to every row of a (batch, n) matrix.
a @ b (or torch.matmul) is matrix multiplication. With 3-d or higher inputs it multiplies the last two dimensions and broadcasts over the leading ones, which is how batched attention is written.
Reductions such as sum, mean, max, argmax, std take a dim argument and remove that dimension; pass keepdim=True to keep it as size 1 so the result still broadcasts against the input. t.max(dim=1) returns both values and indices.
Methods ending in an underscore are in-place: a.add_(1), a.zero_(), a.clamp_(0, 1). They save memory but overwrite data that autograd may need, so avoid them on tensors involved in gradient computation unless you know it is safe.
Everything runs on whatever device the tensors live on, with the same code.
Example
import torch
a = torch.tensor([[1., 2.], [3., 4.]])
b = torch.tensor([10., 20.])
print(a + b) # tensor([[11., 22.], [13., 24.]]) b broadcast to each row
print(a * a) # tensor([[ 1., 4.], [ 9., 16.]]) element-wise
print(a @ a) # tensor([[ 7., 10.], [15., 22.]]) matrix product
print(a.sum()) # tensor(10.)
print(a.sum(dim=0)) # tensor([4., 6.])
print(a.mean(dim=1, keepdim=True).shape) # torch.Size([2, 1])
values, idx = a.max(dim=1)
print(values, idx) # tensor([2., 4.]) tensor([1, 1])
a.add_(1) # in place
print(a) # tensor([[2., 3.], [4., 5.]])Remember
*is element-wise,@is matrix multiplication- Reductions take
dim;keepdim=Truekeeps broadcasting easy - A trailing underscore means the op modifies the tensor in place
Watch out for
- Using
*when you meant a matrix product, which silently broadcasts instead of failing.