diff --git a/aneforge/linalg.py b/aneforge/linalg.py index 3e6c83b..74f15ee 100644 --- a/aneforge/linalg.py +++ b/aneforge/linalg.py @@ -529,11 +529,19 @@ def eigvals(A, iters: int = 60): "eigh", "eigvals", "generalized_eigh", "dominant_eig", "svd", "dominant_svd", "svdvals_topk", "randomized_svd", "pca", "kron", + "polar", ] # __main__ - self-test / validation vs numpy/scipy +def _general_square(n, cond, seed): + """A general (non-symmetric) n x n matrix with the target condition number: U diag(s) V^T.""" + rng = np.random.default_rng(seed) + U = np.linalg.qr(rng.standard_normal((n, n)))[0]; V = np.linalg.qr(rng.standard_normal((n, n)))[0] + return (U * np.geomspace(1.0, cond, n)) @ V.T + + def _make_spd(n, cond, seed): """SPD A with target condition number, fp16-stored (reference solves the fp16-rounded system).""" rng = np.random.default_rng(seed) @@ -791,6 +799,26 @@ def matrix_power(A, n: int): return np.asarray(acc, np.float32) +def polar(A): + """Polar decomposition A = U @ P, on the ANE. + + U is orthogonal for square A and semi-orthogonal for either rectangular shape (U^T U = I for + tall A, U U^T = I for wide A); P is symmetric positive-semidefinite [n,n]. + Composed from the on-ANE SVD: U_ S V^T -> U = U_ V^T, P = V diag(S) V^T. + Oracle: scipy.linalg.polar(A, side='right').""" + A16 = np.asarray(A, f16) + if A16.ndim != 2: raise ValueError(f"polar: expected 2-D; got shape {A16.shape}") + m, n = A16.shape + U_svd, S, Vt = randomized_svd(A16, k=min(m, n), oversample=5, power_iters=2) + V = Vt.T.astype(f16) + Umat = _ane_gemm(U_svd.astype(f16), Vt) # U_ @ V^T -> [m,n] + Sdiag = np.diag(S).astype(f16) + Pmat = _ane_gemm(_ane_gemm(V, Sdiag), Vt) # V diag(S) V^T -> [n,n] + P = np.asarray(Pmat, np.float32) + P = 0.5 * (P + P.T) # fp16 GEMMs leave ~1e-4 asymmetry; sym kills it + return np.asarray(Umat, np.float32), P + + def main(): print("=" * 90) print("aneforge.linalg - ITERATIVE linear algebra on the ANE (matmuls=ANE, RNG/QR/SVD/loop=HOST)") @@ -903,6 +931,23 @@ def main(): f"(~{ane_flops/host_flops:.0f}x on ANE)") print() + # ---------------- polar decomposition ------------------------------- # + print("-" * 90) + print("POLAR DECOMPOSITION (A = U P, U orthogonal, P SPD) - vs scipy.linalg.polar") + print("-" * 90) + import scipy.linalg + for n in (6, 8): + # a GENERAL matrix: for symmetric positive-definite input polar is the trivial case (U = I, P = A) + A_pol = f16(_general_square(n, 1e1, 80 + n)) + U, P = polar(A_pol) + ref_P = scipy.linalg.polar(A_pol, side="right")[1] + recon_err = _relerr(U @ P, A_pol) + orth_err = _relerr(U.T @ U, np.eye(n)) + spd_err = _relerr(P, ref_P) + print(f" n={n}: recon relerr={recon_err:.3e} U^T U ~ I relerr={orth_err:.3e} " + f"P vs scipy relerr={spd_err:.3e}") + print() + # ---------------- verdict ---------------------------------------------- # print("#" * 90) print("# VERDICT - iterative linear algebra on the ANE") diff --git a/tests/test_linalg.py b/tests/test_linalg.py index 496cb5a..7e7c4a4 100644 --- a/tests/test_linalg.py +++ b/tests/test_linalg.py @@ -492,3 +492,52 @@ def test_kron(): max_err = max(max_err, err) assert np.allclose(got16, ref16, atol=5e-4, rtol=0), f"kron({m},{n},{p},{q}) max abs err {err}" assert max_err <= 5e-4 +# ----------------------------- polar decomposition ----------------------------- # + +@pytest.mark.parametrize("n", [6, 8]) +def test_polar_reconstruction(n): + A = f16(_square(n, 1e1, 100 + n)) + U, P = L.polar(A) + assert relerr(U @ P, A) <= 5e-2, f"polar recon: relerr {relerr(U @ P, A)}" + + +@pytest.mark.parametrize("n", [6, 8]) +def test_polar_orthogonality(n): + A = f16(_square(n, 1e1, 110 + n)) + U, _ = L.polar(A) + assert relerr(U.T @ U, np.eye(n)) <= 5e-2, f"U^T U ~ I: relerr {relerr(U.T @ U, np.eye(n))}" + + +@pytest.mark.parametrize("n", [6, 8]) +def test_polar_psd(n): + A = f16(_square(n, 1e1, 120 + n)) + _, P = L.polar(A) + # P should be symmetric (fp16 randomized SVD introduces ~1e-4 asymmetry) + assert relerr(P, P.T) <= 1e-3, f"P not symmetric: relerr {relerr(P, P.T)}" + # P should be PSD (all eigenvalues >= 0) + eig = np.linalg.eigvalsh(P) + assert eig.min() >= -1e-3, f"P has negative eigenvalue: {eig.min()}" + + +def test_polar_matches_scipy(): + scipy_linalg = pytest.importorskip("scipy.linalg") + A = f16(_square(8, 1e1, 130)) + U, P = L.polar(A) + ref_U, ref_P = scipy_linalg.polar(A, side="right") + assert relerr(P, ref_P) <= 5e-2, f"P vs scipy: relerr {relerr(P, ref_P)}" + + +@pytest.mark.parametrize("m,n", [(8, 5), (5, 8)]) +def test_polar_rectangular(m, n): + """Both rectangular shapes: U is [m,n] semi-orthogonal, P is [n,n], and U P reconstructs A.""" + A = np.asarray(np.random.default_rng(3).standard_normal((m, n)), f16) + U, P = L.polar(A) + assert U.shape == (m, n) and P.shape == (n, n) + assert relerr(U @ P, A) <= 5e-2, f"polar recon [{m},{n}]: relerr {relerr(U @ P, A)}" + I = U.T @ U if m >= n else U @ U.T + assert relerr(I, np.eye(I.shape[0])) <= 5e-2, f"U not semi-orthogonal: relerr {relerr(I, np.eye(I.shape[0]))}" + + +def test_polar_rejects(): + with pytest.raises(ValueError): L.polar(np.zeros(4, f16)) # not 2-D + with pytest.raises(ValueError): L.polar(np.zeros((2, 2, 2), f16))