clean svd tests, set full_matrices false in torch backend (#11113)

* clean tests, set full_matrices false

* add more shape asserts
This commit is contained in:
kevvz
2025-07-06 13:55:49 -04:00
committed by GitHub
parent a556f50668
commit b7af9cf849
2 changed files with 9 additions and 8 deletions
+1 -1
View File
@@ -353,7 +353,7 @@ def sort_values(input, dim=-1, descending=False, stable=True, values=None, indic
return wrap(out_values), wrap(out_indices)
@torch.library.impl("aten::_linalg_svd", "privateuseone")
def _linalg_svd(self, full_matrices=True):
def _linalg_svd(self, full_matrices=False):
U, S, Vh = unwrap(self).svd(full_matrices)
return wrap(U), wrap(S), wrap(Vh)
+8 -7
View File
@@ -3061,15 +3061,16 @@ class TestOps(unittest.TestCase):
def test_svd(self):
# test for tiny backend. real svd tests are in test_linalg
A = torch.randn(5, 5)
U, S, Vh = torch.linalg.svd(A, full_matrices=True)
np.testing.assert_allclose(torch.dist(A, U @ torch.diag(S) @ Vh).cpu().numpy(), 0, atol=1e-5)
U, S, Vh = torch.linalg.svd(A, full_matrices=False)
U, S, Vh = torch.linalg.svd(A)
np.testing.assert_equal(U.shape, (5,5))
np.testing.assert_equal(Vh.shape, (5,5))
np.testing.assert_allclose(torch.dist(A, U @ torch.diag(S) @ Vh).cpu().numpy(), 0, atol=1e-5)
# # TODO: this works with torch, but not TINY_BACKEND. U has a wrong shape
# A = torch.randn(5, 3)
# U, S, Vh = torch.linalg.svd(A, full_matrices=False)
# np.testing.assert_allclose(torch.dist(A, U @ torch.diag(S) @ Vh).cpu().numpy(), 0, atol=1e-5)
A = torch.randn(5, 3)
U, S, Vh = torch.linalg.svd(A, full_matrices=False)
np.testing.assert_equal(U.shape, (5,3))
np.testing.assert_equal(Vh.shape, (3,3))
np.testing.assert_allclose(torch.dist(A, U @ torch.diag(S) @ Vh).cpu().numpy(), 0, atol=1e-5)
@unittest.skipUnless(is_dtype_supported(dtypes.uchar), f"no uint8 on {Device.DEFAULT}")
class TestOpsUint8(unittest.TestCase):