mirror of
https://github.com/tinygrad/tinygrad.git
synced 2026-01-20 04:18:13 -05:00
b6835f4134
* remove Ops.VIEW and related UOp methods * update abstractions2.py * no ShapeTrackers in abstractions2.py * it's a size 1
70 lines
2.1 KiB
Python
70 lines
2.1 KiB
Python
import unittest
|
|
from tinygrad import Tensor
|
|
from tinygrad.uop.ops import UPat, Ops, UOp
|
|
|
|
# NOTE: unlike before base for a realized tensor is always a BUFFER
|
|
realized_pattern = UPat(Ops.BUFFER)
|
|
# after realization, base tensor uops become RESHAPE(BUFFER)
|
|
buffer_view_pattern = UPat(Ops.RESHAPE, src=(UPat(Ops.BUFFER),))
|
|
def is_pattern_uop(u:UOp, pat:UPat): assert pat.match(u, {}), f"{u}\nis not\n{pat}"
|
|
def is_pattern(ten:Tensor, pat:UPat): is_pattern_uop(ten.uop, pat)
|
|
|
|
class TestTensorMutates(unittest.TestCase):
|
|
def test_mutate_add(self):
|
|
a = Tensor([1,2,3])
|
|
b = Tensor([4,5,6])
|
|
ret = a+b
|
|
pa = a.uop
|
|
pb = b.uop
|
|
pr = ret.uop
|
|
ret.schedule()
|
|
self.assertIsNot(pa, a.uop)
|
|
self.assertIsNot(pb, b.uop)
|
|
self.assertIsNot(pr, ret.uop)
|
|
for t in [a,b,ret]: is_pattern_uop(t.uop.base, realized_pattern)
|
|
|
|
def test_reshape_is_same_parent(self):
|
|
a = Tensor([1,2,3])
|
|
b = Tensor([4,5,6])
|
|
c = a+b
|
|
d = (a+b).reshape(3,1)
|
|
d.realize()
|
|
is_pattern_uop(d.uop.base, realized_pattern)
|
|
is_pattern_uop(c.uop.base, realized_pattern)
|
|
is_pattern_uop(c.uop.base, realized_pattern)
|
|
assert d.uop is not d.uop.base
|
|
|
|
def test_reshape_is_same_child(self):
|
|
a = Tensor([1,2,3])
|
|
b = Tensor([4,5,6])
|
|
c = a+b
|
|
d = (a+b).reshape(3,1)
|
|
c.realize()
|
|
is_pattern_uop(c.uop.base, realized_pattern)
|
|
is_pattern_uop(d.uop.base, realized_pattern)
|
|
|
|
class TestTensorUopRepresentation(unittest.TestCase):
|
|
def test_realized(self):
|
|
a = Tensor([1.,2,3]).realize()
|
|
print(a.uop)
|
|
is_pattern_uop(a.uop.base, realized_pattern)
|
|
|
|
def test_add_realized(self):
|
|
a = Tensor([1.,2,3]).realize()
|
|
b = Tensor([4.,5,6]).realize()
|
|
c = a+b
|
|
print(c.uop)
|
|
is_pattern(c, UPat(Ops.ADD))
|
|
for s in c.uop.src: is_pattern_uop(s.base, realized_pattern)
|
|
|
|
def test_empty_buf(self):
|
|
a = Tensor.empty(3, 3)
|
|
is_pattern(a, UPat(Ops.RESHAPE, src=(UPat(Ops.BUFFER),)))
|
|
vi = UOp.variable("i", 1, 3).bind(1)
|
|
a = Tensor.empty(3, vi)
|
|
is_pattern(a, UPat(Ops.RESHAPE, src=(UPat(Ops.SHRINK, src=(UPat(Ops.BUFFER),))),))
|
|
self.assertEqual(a.uop.base.buffer.size, 9)
|
|
|
|
if __name__ == '__main__':
|
|
unittest.main()
|