Tensor Math: Element-wise Ops, Matmul and Reductions

Almost every neural network is built from three kinds of tensor operations: element-wise operations that act on each number separately, matrix multiplication that mixes values along a dimension, and reductions that collapse a dimension into a summary such as a sum or mean. This post covers all three with runnable examples.

Element-wise operations

Arithmetic operators work on matching positions and return a new tensor.

import torch

a = torch.tensor([1.0, 2.0, 3.0])
b = torch.tensor([4.0, 5.0, 6.0])

print(a + b)          # tensor([5., 7., 9.])
print(a * b)          # tensor([ 4., 10., 18.])  element-wise, not dot product
print(a ** 2)         # tensor([1., 4., 9.])
print(a / b)
print(torch.exp(a))
print(torch.sqrt(a))
print(torch.clamp(a, min=1.5, max=2.5))   # tensor([1.5, 2.0, 2.5])

Mathematical functions such as exp, log, sqrt, sin, tanh and sigmoid are all element-wise. A scalar combines with every element, so a * 2 doubles each value.

In-place operations

Methods ending in an underscore modify the tensor in place, for example a.add_(1) or a.zero_(). They save memory but can break autograd if they overwrite a value needed for the backward pass, so prefer the normal versions in model code.

Matrix multiplication

A = torch.rand(2, 3)
B = torch.rand(3, 4)

C = A @ B                 # same as torch.matmul(A, B)
print(C.shape)            # torch.Size([2, 4])

v = torch.tensor([1.0, 2.0, 3.0])
w = torch.tensor([4.0, 5.0, 6.0])
print(v @ w)              # dot product: tensor(32.)
print(torch.dot(v, w))    # same

For matrices the inner dimensions must match: (2, 3) @ (3, 4) gives (2, 4). A linear layer is exactly this operation plus a bias. matmul also handles batches: with shapes (8, 2, 3) @ (8, 3, 4) it multiplies each pair and returns (8, 2, 4). This batched form is the basis of attention in transformers.

x = torch.rand(8, 2, 3)
y = torch.rand(8, 3, 4)
print(torch.bmm(x, y).shape)       # torch.Size([8, 2, 4])
print((x @ y).shape)               # same result
print(A.T.shape)                   # transpose of a 2-D tensor: (3, 2)

Reductions

A reduction collapses one or more dimensions. Without a dim argument it reduces everything to a single number.

x = torch.tensor([[1.0, 2.0, 3.0],
                  [4.0, 5.0, 6.0]])

print(x.sum())                    # tensor(21.)
print(x.sum(dim=0))               # tensor([5., 7., 9.])  per column
print(x.sum(dim=1))               # tensor([ 6., 15.])    per row
print(x.mean(dim=1))              # tensor([2., 5.])
print(x.sum(dim=1, keepdim=True)) # shape (2, 1) instead of (2,)
print(x.max())                    # tensor(6.)

Remember: dim is the dimension that disappears. Summing dim=0 on a (2, 3) tensor leaves shape (3,). Use keepdim=True when you want to divide the original tensor by the result, because the kept size-one dimension broadcasts correctly.

max, argmax and softmax

scores = torch.tensor([[0.1, 2.0, 0.3],
                       [1.5, 0.2, 0.9]])

values, indices = scores.max(dim=1)
print(values)                     # tensor([2.0000, 1.5000])
print(indices)                    # tensor([1, 0])
print(scores.argmax(dim=1))       # tensor([1, 0]): predicted classes

probs = torch.softmax(scores, dim=1)
print(probs.sum(dim=1))           # tensor([1., 1.])

When you call max(dim=...) you get both the largest values and where they were. In classification, argmax(dim=1) on the output scores gives the predicted class for each sample. softmax turns scores into probabilities that sum to one along the chosen dimension.

Thinking in shapes

When you write tensor math, think about shapes first and numbers second. A fully connected layer maps an input of shape (batch, in_features) to (batch, out_features) by multiplying with a weight matrix of shape (in_features, out_features), then adding a bias vector. In PyTorch, nn.Linear stores the weight transposed, with shape (out_features, in_features), and applies it as x @ weight.T + bias. Writing this by hand once is a useful exercise: create random input and weights, multiply, add the bias, and compare the result with the output of an actual nn.Linear layer after copying its weights.

Reductions give you the building blocks of loss functions and metrics. The mean squared error is simply ((pred - target) ** 2).mean(). Classification accuracy is (pred.argmax(dim=1) == labels).float().mean(): compare predicted and true classes to get a bool tensor, convert it to float, and average. Each of these lines is a short pipeline of element-wise operations followed by a reduction, which is why understanding the basic operations pays off so quickly.

Finally, be aware of numerical stability. Computing torch.log(torch.softmax(x, dim=1)) can produce infinities for large inputs, so PyTorch provides torch.log_softmax and fused loss functions that compute the same quantity safely. Prefer these built-in combinations over writing the formula yourself.

Common mistakes

• Using * when you mean matrix multiplication. Use @ or matmul.
• Choosing the wrong dim. Check the output shape after every reduction.
• Integer tensors: mean fails on integers. Convert with .float() first.
• Applying softmax on the wrong dimension, which silently gives wrong probabilities.

Next, broadcasting explains how PyTorch combines tensors of different shapes in element-wise operations.