github/tinygrad
1
1
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-04-29 03:00:14 -04:00
Code Issues Packages Projects Releases Wiki Activity

scheduler ops_folding from delete_lazy (#8124)

Browse Source 
* scheduler diff from delete_lazy

* test_std_mean

* late fold copy of CONST

* clang const is fine
This commit is contained in:
qazal
2024-12-09 18:36:01 +02:00
committed by GitHub
parent 87c360c4b5
commit 80de06c8b9
4 changed files with 30 additions and 13 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
test/test_ops.py
View File

@@ -1224,10 +1224,13 @@ class TestOps(unittest.TestCase):
lambda x: Tensor.stack(*x.std_mean(correction=5)))
helper_test_op([(15,25,35)], lambda x: torch.stack(torch.std_mean(x, keepdim=True, correction=0)),
lambda x: Tensor.stack(*x.std_mean(keepdim=True, correction=0)))
helper_test_op([(1,0,3,0,5)], lambda x: torch.stack(torch.std_mean(x, axis=(1,3))),
lambda x: Tensor.stack(*x.std_mean(axis=(1,3))))
helper_test_op([(3,4,5,6)], lambda x: torch.stack(torch.std_mean(x, axis=(1,2))),
lambda x: Tensor.stack(*x.std_mean(axis=(1,2))))
@unittest.skip("TODO: this fails because of loaded nan in mul folding")
def test_std_mean_loaded_nan(self):
helper_test_op([(1,0,3,0,5)], lambda x: torch.stack(torch.std_mean(x, axis=(1,3))),
lambda x: Tensor.stack(*x.std_mean(axis=(1,3))))
def test_softmax(self):
helper_test_op([(45,65)], torch.nn.Softmax(dim=1), Tensor.softmax, atol=1e-7, grad_atol=1e-7)
helper_test_op([(45)], torch.nn.Softmax(dim=0), Tensor.softmax, atol=1e-7, grad_atol=1e-7)