mawei/tinygrad
1
0
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-06-13 00:15:35 +08:00
Code Issues Packages Projects Releases Wiki Activity

Implement linalg_eigh function for tensor eigenvalue decomposition in torch backend (#10612)

Browse Source 
* Implement private _linalg_eigh function for tensor eigenvalue decomposition in torch backend

* Add unit test for linalg.eigh function in TestTorchBackend

This test verifies the eigenvalue decomposition of a 2x2 tensor using the linalg.eigh function, ensuring the computed eigenvalues and reconstructed tensor match the expected results.
This commit is contained in:
Xingyu
2025-06-04 19:59:50 +08:00
committed by GitHub
parent 7114b6ab31
commit 7a1bfb668d
2 changed files with 13 additions and 0 deletions
Download Patch File Download Diff File Expand all files Collapse all files

6
extra/torch_backend/backend.py
View File

@@ -121,6 +121,12 @@ def cummax(self, dim):
# TODO: move to tinygrad
def nonzero(self): return aten.nonzero(self.cpu()).tiny()
@torch.library.impl("aten::_linalg_eigh", "privateuseone")
# TODO: move to tinygrad
def _linalg_eigh(self, UPLO: str = 'U'):
w, v = torch.linalg.eigh(self.cpu(), UPLO=UPLO)
return w.tiny(), v.tiny()
def upsample_backward(grad_out, output_size, input_size, *args, f=None): return f(grad_out.cpu(), output_size, input_size, *args).tiny()
for i in [

7
extra/torch_backend/test.py
View File

@@ -170,6 +170,13 @@ class TestTorchBackend(unittest.TestCase):
assert torch.equal(tensor_a, tensor_b)
assert not torch.equal(tensor_a, tensor_c)
def test_linalg_eigh(self):
a = torch.tensor([[1, 2], [2, 1]], dtype=torch.float32, device=device)
w, v = torch.linalg.eigh(a)
np.testing.assert_equal(w.cpu().numpy(), [-1, 3])
recon = (v @ torch.diag(w) @ v.T).cpu().numpy()
np.testing.assert_allclose(recon, a.cpu().numpy(), atol=1e-6)
def test_scalar_assign(self):
a = torch.tensor([1, 2, 3], device=device)
a[1] = 4