torch.linalg

Active contributors: lezcano, IvanYashchuk, nikitaved

Purpose

torch.linalg is PyTorch's NumPy-compatible numerical linear algebra module. It exposes factorizations (LU, QR, Cholesky, eigen, SVD), solves, norms, determinants, and matrix functions. Most users come here from numpy.linalg and find a one-to-one analog.

Surface area

import torch
A = torch.randn(3, 3)
torch.linalg.cholesky(A @ A.T)
torch.linalg.qr(A)
torch.linalg.svd(A)
torch.linalg.solve(A, b)
torch.linalg.eig(A)
torch.linalg.norm(A, ord="fro")
torch.linalg.matrix_rank(A)
torch.linalg.pinv(A)
torch.linalg.det(A)
torch.linalg.lstsq(A, b)

The full list lives in torch/linalg/__init__.py.

How the kernels work

Most decompositions delegate to vendor libraries:

• CPU: LAPACK (Intel MKL, OpenBLAS, NVPL) via aten/src/ATen/native/BatchLinearAlgebra.cpp.
• CUDA: cuSolver and MAGMA via aten/src/ATen/native/cuda/linalg/.
• ROCm: rocBLAS / rocSOLVER via the hipified CUDA path.

at::globalContext().linalgPreferredBackend() (and the env var TORCH_LINALG_PREFER_*) lets users switch between cuSolver and MAGMA at runtime.

Some smaller routines (norms, transpositions, basic matmul) live in aten/src/ATen/native/LinearAlgebra.cpp and run as ordinary CPU/CUDA kernels.

Differentiability

Almost every routine is differentiable. The closed-form gradients are in tools/autograd/derivatives.yaml and the helper functions in aten/src/ATen/native/BatchLinearAlgebraKernel.cpp. SVD, eig, and Cholesky have particularly subtle gradient formulas; the implementations have been the subject of multiple research-paper-grade refactors.

Batched semantics

Every routine works on batched inputs of shape [..., M, N]. The implementations call into batched LAPACK/cuSolver entry points; for shapes those libraries don't support, PyTorch falls back to a loop.

Specific niches

• Tridiagonal solve. torch.linalg.solve_triangular and solve_ex for triangular systems.
• Generalized eigenvalue problem. torch.linalg.eigh for Hermitian / symmetric.
• Multi-dot. torch.linalg.multi_dot chains matmuls with optimal parenthesization.
• Matrix exponential. torch.linalg.matrix_exp.
• Tensor solve / contraction. torch.linalg.tensorsolve, torch.linalg.tensorinv.

NumPy compatibility

torch.linalg follows the NumPy API spec where possible. The two main wrinkles:

• PyTorch's complex support is more recent; some routines in older versions only worked on real inputs.
• Backward of certain routines requires constraints (e.g., SVD backward requires distinct singular values for non-degenerate gradients).

Where to look

File	Purpose
torch/linalg/__init__.py	Public surface
aten/src/ATen/native/LinearAlgebra.cpp	Basic routines (norms, dot, trace)
aten/src/ATen/native/BatchLinearAlgebra.cpp	LAPACK-backed batched decompositions
aten/src/ATen/native/cuda/linalg/	cuSolver/MAGMA-backed CUDA
tools/autograd/derivatives.yaml (linalg entries)	Backward formulas

Where to read next

• Systems / ATen — kernel mechanics.
• Features / Mixed precision — TF32 / mixed-precision matmul.