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

simpler Tensor.pow to integer (#2746)

Browse Source
This commit is contained in:
chenyu
2023-12-13 11:39:20 -05:00
committed by GitHub
parent 26f49869f4
commit aa4a0de287
3 changed files with 12 additions and 5 deletions
Download Patch File Download Diff File Expand all files Collapse all files

9
test/test_schedule.py
View File

@@ -327,11 +327,18 @@ class TestSchedule(unittest.TestCase):
out = x.to('cpu')
check_schedule(out, 0, filter_loadops=False)
def test_pow_const_tensor(self):
def test_pow_const_tensor_simplified(self):
x = Tensor([1,2,3,4])
# NOTE: this does not test ** Tensor(2) is simpler in ast than ** Tensor(2.5)
out = x ** Tensor(2)
check_schedule(out, 1)
def test_pow_const_tensor_to_zero(self):
x = Tensor([1,2,3,4])
out = x ** Tensor(0)
# NOTE: this is ConstBuffer 0 + ConstBuffer 1
check_schedule(out, 1)
def test_zero_size(self):
x = Tensor.rand(2, 3, 0)
out = x + 1