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
| Function | Description |
|---|---|
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
- Tensors are compared from the trailing dimension.
- Dimensions of size 1 are expanded to match the other.
- 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