Turbocharging SVD with JAX

In the previous post, I wrote about the fundamentals of two commonly used dimensionality reduction approaches, singular value decomposition (SVD) and principal component analysis (PCA). I also explained their relationships using numpy. To quickly recap, the singular values (Σ) of a 0-centered matrix X (n samples × m features), equals the square root of its eigenvalues (Λ), making it possible to compute PCA using SVD, which is often more efficient.

In this post, I will explore and benchmark a few SVD algorithms and implementations in Python for squared dense matrix (X, n×n). I also…