PyTorch Tensor Operations

Tensor operations are the math instructions that PyTorch executes when it processes data or trains a model. Knowing which operation to use and when makes the difference between writing clear, efficient code and spending hours debugging shape errors.

Categories of Tensor Operations

Tensor Operations
├── Arithmetic        (add, subtract, multiply, divide)
├── Comparison        (greater than, equal to, less than)
├── Reduction         (sum, mean, max, min)
├── Shape             (reshape, squeeze, unsqueeze, transpose)
├── Joining           (cat, stack)
└── Linear Algebra    (matmul, inverse, norm)

Arithmetic Operations

All arithmetic operations work element-wise — meaning PyTorch applies the operation to matching positions across two tensors.

import torch

a = torch.tensor([10.0, 20.0, 30.0])
b = torch.tensor([2.0, 4.0, 5.0])

print(torch.add(a, b))       # tensor([12., 24., 35.])
print(torch.subtract(a, b))  # tensor([ 8., 16., 25.])
print(torch.multiply(a, b))  # tensor([ 20.,  80., 150.])
print(torch.divide(a, b))    # tensor([5., 5., 6.])

You can also use Python operators directly: a + b, a - b, a * b, a / b.

In-Place Operations

In-place operations modify the tensor directly without creating a new one. They save memory. In PyTorch, in-place operations end with an underscore.

t = torch.tensor([1.0, 2.0, 3.0])
t.add_(10)     # adds 10 to every element IN PLACE
print(t)       # tensor([11., 12., 13.])

Use in-place operations carefully during training — they can interfere with gradient computation if applied to tensors that require gradients.

Comparison Operations

Comparison operations return boolean tensors — each position is True or False depending on whether the condition holds.

a = torch.tensor([1.0, 5.0, 3.0, 7.0, 2.0])

print(a > 3)           # tensor([False,  True, False,  True, False])
print(a == 3)          # tensor([False, False,  True, False, False])
print(torch.equal(a, a))   # True — checks if two tensors are identical

Reduction Operations

Reduction operations collapse a tensor along one or more dimensions, producing a smaller tensor.

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

print(t.sum())           # tensor(21.)  ← sum of all elements
print(t.sum(dim=0))      # tensor([5., 7., 9.])  ← column sums
print(t.sum(dim=1))      # tensor([ 6., 15.])  ← row sums
print(t.mean())          # tensor(3.5)
print(t.max())           # tensor(6.)

dim=0 means reduce across rows (down columns):
  [1, 2, 3]       col sum:  1+4=5, 2+5=7, 3+6=9
+ [4, 5, 6]
= [5, 7, 9]

dim=1 means reduce across columns (across rows):
  [1, 2, 3]  → row sum:  1+2+3=6
  [4, 5, 6]  → row sum:  4+5+6=15

Shape Operations

squeeze and unsqueeze

squeeze removes dimensions of size 1. unsqueeze adds a dimension of size 1. These are used constantly to make tensors compatible for operations that expect a specific number of dimensions.

t = torch.tensor([1.0, 2.0, 3.0])
print(t.shape)               # torch.Size([3])

t2 = t.unsqueeze(0)          # adds batch dimension at position 0
print(t2.shape)              # torch.Size([1, 3])

t3 = t2.squeeze(0)           # removes the size-1 dimension
print(t3.shape)              # torch.Size([3])

unsqueeze(0):
[1, 2, 3]   →   [[1, 2, 3]]
shape (3,)       shape (1, 3)
                 └── new dimension added here

transpose and permute

transpose swaps two dimensions. permute reorders all dimensions.

t = torch.rand(3, 4)      # shape: 3 rows, 4 columns
print(t.T.shape)           # torch.Size([4, 3]) ← transposed

# permute: rearrange a 3D tensor from (C, H, W) to (H, W, C)
img = torch.rand(3, 224, 224)         # channels, height, width
img_hwc = img.permute(1, 2, 0)        # height, width, channels
print(img_hwc.shape)   # torch.Size([224, 224, 3])

Joining Tensors

cat – Concatenate Along an Existing Dimension

a = torch.tensor([[1.0, 2.0],
                  [3.0, 4.0]])   # shape: 2x2

b = torch.tensor([[5.0, 6.0]])   # shape: 1x2

c = torch.cat([a, b], dim=0)    # join along rows
print(c)
# tensor([[1., 2.],
#         [3., 4.],
#         [5., 6.]])   shape: 3x2

stack – Create a New Dimension

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

z = torch.stack([x, y], dim=0)   # shape: 2x3
print(z)
# tensor([[1., 2., 3.],
#         [4., 5., 6.]])

cat vs stack:
  cat: does NOT add new dimensions — glues tensors together along existing axis
  stack: ADDS a new dimension — treats each tensor as one element in a new axis

Math Functions

t = torch.tensor([1.0, 4.0, 9.0, 16.0])

print(torch.sqrt(t))      # tensor([1., 2., 3., 4.])
print(torch.exp(t))       # exponential: e^x for each element
print(torch.log(t))       # natural log
print(torch.abs(torch.tensor([-1.0, 2.0, -3.0])))   # absolute value

Summary

PyTorch tensor operations fall into six categories: arithmetic, comparison, reduction, shape manipulation, joining, and linear algebra. Arithmetic works element-wise. Reduction collapses tensors along specified dimensions. squeeze and unsqueeze manage extra dimensions. cat joins tensors along existing dimensions while stack creates a new one. Understanding these operations removes the confusion from most shape-related errors you encounter while building models.