github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-01-09 15:08:02 -05:00
Code Issues Packages Projects Releases Wiki Activity

Fix SVD rank‑1 Jacobi rotation when tau == 0 (#13945)

Browse Source
This commit is contained in:
haofei
2025-12-31 21:30:18 -08:00
committed by GitHub
parent 20777f30b9
commit 526fd4ec71
2 changed files with 7 additions and 1 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
test/unit/test_linalg.py
View File

@@ -83,6 +83,12 @@ class TestLinAlg(unittest.TestCase):
s_diag = (S.unsqueeze(-2) * Tensor.eye(2))
reconstruction_helper([U, s_diag, V], a)
def test_svd_rank1(self):
a = Tensor([[1.0, 1.0], [2.0, 2.0]]).realize()
U, S, V = a.svd()
np.testing.assert_allclose(S.numpy(), [np.sqrt(10), 0.0], atol=1e-4, rtol=1e-4)
reconstruction_helper([U, S.unsqueeze(-2) * Tensor.eye(2), V], a)
def test_newton_schulz(self):
coefficients = [(2, -1.5, 0.5), (2.0, -1.4, 0.2, 0.2)]#these params map to the sign function
sizes = [(2,2), (3,2), (2,3), (2,2,2)]

2
tinygrad/tensor.py
View File

@@ -3672,7 +3672,7 @@ class Tensor(OpMixin):
alpha, beta = U_permuted.square().sum(-2).unsqueeze(-2).split(num//2, -1)
rot = gamma != 0
tau = (beta - alpha) / (2 * rot.where(gamma, 1))
t = tau.sign() / (tau.abs() + (1 + tau.square()).sqrt())
t = (tau != 0).where(tau.sign(), 1) / (tau.abs() + (1 + tau.square()).sqrt())
t = rot.where(t, 0)
c = 1 / (1 + t.square()).sqrt()
s = c * t