PyTorch Cheatsheet

Tensor Operations

Use this PyTorch reference while you build software engineering projects, review code for technical interview prep, or polish examples for a software engineer resume.

Arithmetic

# Element-wise — functional and operator forms both work
a + b          # torch.add(a, b)
a - b          # torch.sub(a, b)
a * b          # torch.mul(a, b)
a / b          # torch.div(a, b)
a // b         # torch.div(a, b, rounding_mode='floor')
a % b          # torch.remainder(a, b)
a ** b         # torch.pow(a, b)
-a             # torch.neg(a)

# In-place (trailing underscore)
a.add_(b)
a.mul_(2)

In-place ops on tensors that require grad may cause errors during backprop. Prefer out-of-place during training.

Matrix Operations

a @ b                        # matrix multiply (matmul)
torch.matmul(a, b)           # same — handles batches and broadcasting
torch.mm(a, b)               # 2-D only, no broadcast
torch.bmm(a, b)              # batched matmul (a: B×M×K, b: B×K×N)
torch.mv(matrix, vec)        # matrix-vector product
torch.dot(a, b)              # 1-D dot product (no broadcast)
torch.outer(a, b)            # outer product

torch.linalg.solve(A, b)     # Ax = b
torch.linalg.lstsq(A, b)     # least squares
torch.linalg.inv(A)          # matrix inverse
torch.linalg.det(A)          # determinant
torch.linalg.matrix_rank(A)
torch.linalg.norm(t)         # Frobenius by default
torch.linalg.svd(A)          # singular value decomposition → U, S, Vh
torch.linalg.eig(A)          # eigendecomposition
torch.linalg.cholesky(A)     # Cholesky decomposition
torch.linalg.qr(A)           # QR decomposition

Reduction Operations

FunctionDescription
t.sum() / torch.sum(t)sum of all elements
t.sum(dim=1)sum along dim 1
t.sum(dim=1, keepdim=True)keep collapsed dim as size-1
t.mean()mean (requires float)
t.std()standard deviation
t.var()variance
t.min() / t.max()scalar min/max
t.min(dim=0)namedtuple (values, indices)
t.argmin() / t.argmax()index of min/max
t.prod()product of all elements
t.median()median
t.mode()mode (most frequent value)
t.norm(p=2)L2 norm (or Lp for other p)
t.cumsum(dim=0)cumulative sum
t.cumprod(dim=0)cumulative product
t.any()True if any element is True
t.all()True if all elements are True
t.count_nonzero()count of non-zero elements
t.sum(dim=(0, 2), keepdim=True)   # sum over multiple dims at once
torch.amax(t, dim=-1)              # max without returning indices
torch.amin(t, dim=-1)
torch.aminmax(t, dim=0)            # both min and max in one call

Element-wise Math

torch.abs(t)          # |x|
torch.sqrt(t)         # √x
torch.rsqrt(t)        # 1/√x (fast)
torch.exp(t)          # eˣ
torch.log(t)          # ln(x)
torch.log2(t)
torch.log10(t)
torch.log1p(t)        # ln(1+x) — numerically stable for small x
torch.expm1(t)        # eˣ − 1
torch.sin(t) / torch.cos(t) / torch.tan(t)
torch.asin(t) / torch.acos(t) / torch.atan(t)
torch.atan2(y, x)
torch.sinh(t) / torch.cosh(t) / torch.tanh(t)
torch.ceil(t) / torch.floor(t) / torch.round(t)
torch.trunc(t)        # truncate toward zero
torch.frac(t)         # fractional part
torch.sign(t)         # -1, 0, or 1
torch.clamp(t, min=0, max=1)    # clip values
torch.clamp_min(t, 0)
torch.clamp_max(t, 1)
torch.sigmoid(t)      # 1/(1+e⁻ˣ)
torch.relu(t)         # max(0, x)  — functional, no param

Comparison Operations

t == other       # torch.eq(t, other)
t != other       # torch.ne
t >  other       # torch.gt
t >= other       # torch.ge
t <  other       # torch.lt
t <= other       # torch.le

torch.isnan(t)
torch.isinf(t)
torch.isfinite(t)
torch.isclose(a, b, rtol=1e-5, atol=1e-8)
torch.allclose(a, b)   # bool — useful in tests
torch.equal(a, b)      # exact element-wise equality (bool)

Sorting and Searching

t.sort(dim=0)                # namedtuple (values, indices)
t.sort(dim=0, descending=True)
t.argsort(dim=0)             # indices that would sort t

torch.topk(t, k=3)          # top-k (values, indices)
torch.topk(t, k=3, largest=False)  # bottom-k

torch.kthvalue(t, k, dim=0)  # k-th smallest (values, indices)

torch.nonzero(t)             # indices of non-zero elements, shape (N, ndim)
torch.nonzero(t, as_tuple=True)   # tuple of 1-D index tensors
torch.where(condition)       # alias for nonzero when 1 arg given
torch.searchsorted(sorted_t, values)

Combining Tensors

torch.cat([a, b, c], dim=0)       # concatenate along existing dim
torch.stack([a, b, c], dim=0)     # stack along NEW dim (all same shape)
torch.hstack([a, b])              # horizontal stack (dim 1)
torch.vstack([a, b])              # vertical stack (dim 0)
torch.dstack([a, b])              # depth stack (dim 2)
torch.block_diag(a, b)            # block-diagonal matrix

Splitting Tensors

torch.split(t, split_size=2, dim=0)    # list of equal chunks
torch.split(t, [1, 2, 3], dim=0)       # variable-size chunks
torch.chunk(t, chunks=3, dim=0)        # up to N equal pieces
torch.unbind(t, dim=0)                 # tuple of slices (drops that dim)

Broadcasting Rules

  1. Tensors are compared from the trailing dimension.
  2. Dimensions of size 1 are expanded to match the other.
  3. Dimensions of size 1 can be on either side.
a = torch.ones(3, 1)   # shape (3, 1)
b = torch.ones(1, 4)   # shape (1, 4)
(a + b).shape           # → (3, 4)

# Explicitly expand (no data copy)
a.expand(3, 4)
a.expand_as(b)          # expand to match b's shape

Random Sampling

torch.manual_seed(42)              # reproducibility

torch.rand(3, 3)                   # uniform [0, 1)
torch.randn(3, 3)                  # standard normal
torch.randint(0, 10, (3, 3))       # ints in [0, 10)
torch.randperm(10)                 # random permutation of 0..9
torch.bernoulli(prob_tensor)       # draw Bernoulli samples

# Distribution objects
dist = torch.distributions.Normal(loc=0.0, scale=1.0)
dist.sample((5,))
dist.log_prob(t)

Scatter and Gather

# gather: pick elements along dim using index tensor
# out[i][j] = input[i][index[i][j]] for dim=1
out = t.gather(dim=1, index=idx)

# scatter_: write values into self at positions in index
t.scatter_(dim=1, index=idx, src=values)
t.scatter_(dim=1, index=idx, value=0.0)  # fill with scalar

# scatter_add_: accumulate
t.scatter_add_(dim=0, index=idx, src=values)

# index_select, index_add_
t.index_select(dim=0, index=torch.tensor([0, 2]))
t.index_add_(dim=0, index=idx, source=src)

Einsum

# Einstein summation — extremely general
torch.einsum('ij,jk->ik', a, b)         # matmul
torch.einsum('bik,bkj->bij', a, b)      # batched matmul
torch.einsum('ij->j', a)                # column sum
torch.einsum('ii->', a)                 # trace
torch.einsum('ij,ij->', a, b)           # Frobenius dot product
torch.einsum('bi,bo->bio', x, y)        # outer product per batch

Type Casting and Conversions

t.float()          # → float32
t.double()         # → float64
t.half()           # → float16
t.long()           # → int64
t.int()            # → int32
t.bool()           # → bool
t.to(dtype=torch.bfloat16)
torch.result_type(a, b)   # what dtype an op would produce