Implementing SVD in tinygrad

The SVD Decomposition

For good reasons every major library NumPy, Scikit-learn, SciPy, PyTorch, Tensorflow, JAX, pandas relies on LAPACK/openBLAS/magma/cuSOLVER/MKL/... implementations of SVD.

See my Tinygrad PR #10813 for more details.

Here is the function that we will implement, part of Tinygrad's Tensor class:

def svd(self:Tensor, k:int|None=None) -> tuple[Tensor, Tensor, Tensor]:
    """
    Returns the SVD decomposition U, S, Vt such that A ≈ U @ diag(S) @ Vt.
    k: Number of singular values to compute (default: min(m,n))
    ```python exec="true" source="above" session="tensor" result="python"
    A = Tensor.randn(5, 10)
    U, S, Vt = A.svd()
    print(U.numpy(), S.numpy(), Vt.numpy())
    ```
    """

In general we want predictable deterministic algorithms, and unfortunately the SVD of a matrix is not unique, even for full rank matrices with sorted distinct singular values we have $2^{\min(m,n)}$ degrees of freedom. See Table 1 for a summary.

Hence, all library sorts the singular values in descending order, which also makes sense because the largest singular value is often the most sought after.

Table 1. Enumeration of the size of the invariant sub-spaces of SVD.

Case	Rank	Singular values	Freedom in $U$ and $V$
1. Full rank, all singular values distinct	$r = \min(m, n)$	$\sigma_1 > \cdots > \sigma_r > 0$	Sign flip on each matching column in $U$ and $V$
2. Full rank, some singular values repeated	$r = \min(m, n)$	Some $\sigma_i = \sigma_{i+1} > 0$	Any rotation/reflection within each degenerate block of columns, applied identically to $U$ and $V$
3. Rank-deficient, all nonzero singular values distinct	$r < \min(m, n)$	$\sigma_1 > \cdots > \sigma_r > 0$, $\sigma_{r+1} = \cdots = 0$	First $r$ columns: sign flip in $U$ and $V$. Remaining: arbitrary rotation in left null space ($U$: $m - r$ cols) and in right null space ($V$: $n - r$ cols), independently
4. Rank-deficient, some positive singular values repeated	$r < \min(m, n)$	Some equal $\sigma_i > 0$ among first $r$, rest zeros	Repeated positive blocks: same rotation in $U$ and $V$. Null spaces: arbitrary independent rotations ($U$: $m - r$ cols, $V$: $n - r$ cols)
5. Zero matrix	$r = 0$	All $\sigma_i = 0$	Arbitrary rotation in both $U$ and $V$.

The implementation

For any matrix $A \in \mathbb{R}^{m \times n}$, the SVD factorizes it as:

$$A = U \Sigma V^T$$

where:

$U \in \mathbb{R}^{m \times \min(m,n)}$ is a matrix with orthonormal columns
$\Sigma \in \mathbb{R}^{\min(m,n) \times \min(m,n)}$ is a diagonal matrix of singular values
$V \in \mathbb{R}^{n \times \min(m,n)}$ is a matrix with orthonormal columns

It turns out that repeated multiplcation of matrix converges to the eigenvectors of the matrix. Here is the full implementation. Also, the following code might seem slightly code-golfed, but this is price you pay for geohotz' line counting.

def svd(self:Tensor, k:int|None=None) -> tuple[Tensor, Tensor, Tensor]:
    # Handle arbitrary batch dims
    *b, m, n = self.shape
    A = self.reshape(-1, m, n)
    if k is None: k = min(m, n)
    def power_iteration(M:Tensor, k:int, max_iter:int=20) -> Tensor:
        B, n, _ = M.shape
        Q, _ = Tensor.randn(B, n, k).qr()
        for _ in range(max_iter): Q, _ = (M @ Q).qr()
        return Q
    # Assume A is tall
    if m < n: A = A.transpose(-2, -1)
    V = power_iteration(A.transpose(-2, -1) @ A, k)
    AV = A @ V
    S = (AV * AV).sum(axis=-2).sqrt().clamp(min_=1e-12)
    U = AV / S.unsqueeze(-2)
    if m < n: U, V = V, U
    # Sort singular values in descending order
    S, idx = S.sort(dim=-1, descending=True)
    U = U.gather(-1, idx.unsqueeze(-2).expand(-1, m, -1))
    V = V.gather(-1, idx.unsqueeze(-2).expand(-1, n, -1))
    return U.reshape(*b, m, k), S.reshape(*b, k), V.reshape(*b, n, k).transpose(-2, -1)

A neat trick we used here to only consider tall matrices is to transpose the matrix if it is wide, since the SVD decomposition of a matrix's transpose just swaps the $U$ and $V$ matrices.

if m < n: A = A.transpose(-2, -1)
U, S, V = A.svd()
if m < n: U, V = V, U

The Condition Number

Our methods square the condition numbers due to the gram matrix $A^T A$. I.e if $c(A)=\kappa$, $c(A^T A)=c(A)^2=\kappa^2$

Test cases

1. Reconstruction Property

The fundamental property: $A = U \Sigma V^T$ or

recon = U @ torch.diag_embed(S) @ V
np.testing.assert_allclose(A.numpy(), recon.numpy(), rtol=1e-3, atol=1e-3)

is trivially satisfied in our case since the method creates $U$ such that it's satisfied.

2. Orthogonality Conditions

The matrices $U$ and $V$ must have orthonormal columns:

$$U^T U = I_{\min(m,n)}$$ $$V V^T = I_{\min(m,n)}$$

I.e

np.testing.assert_allclose((U.mT @ U).numpy(), torch.eye(min(m, n)).expand(*b, -1, -1).numpy(), atol=1e-1)
np.testing.assert_allclose((V @ V.mT).numpy(), torch.eye(min(m, n)).expand(*b, -1, -1).numpy(), atol=1e-1)

Note the different transposes: $U^T U$ for tall/square matrices and $V V^T$ for wide/square matrices.

3. Singular Values Consistency with torch's linalg api

The singular values should match reference implementations:

_, S_torch, _ = torch.linalg.svd(At, full_matrices=False)
np.testing.assert_allclose(St.numpy(), S_torch.numpy(), rtol=1e-2, atol=1e-1)

Why not return full matrices like `torch.linalg.svd`?

Gradients are undefined anyway for the full matrix case
Memory usage is higher
An additional algorithm to implement

PyTorch's SVD Strategy

PyTorch's torch.linalg.svd uses a dispatch strategy based on matrix properties:

General matrices: 'gesvdj' (Jacobi method) - iterative eigenvalue-like approach
Tall or wide matrices (m >> n or m << n): 'gesvda' (Approximate method)
Default fallback: 'gesvd' (QR-based method) - more stable but slower

The Jacobi Alternative

https://github.com/tinygrad/tinygrad/pull/10875

Performance Optimization: tinygrad.jit

Like in the Jacobi version, tinygrad.jit significantly speeds up the QR decomposition.