tinygrad/tinygrad
2
0
Fork 3
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-08-29 17:36:07 +00:00
Code Issues Packages Projects Releases Wiki Activity

don't use view in copy [pr] (#8704)

Browse Source 
* don't use view in copy [pr]

* oh, remove double contig

* fix reps
This commit is contained in:
George Hotz
2025-01-21 09:57:47 -08:00
committed by GitHub
parent d6bf1feaab
commit 018edd934b
3 changed files with 13 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/unit/test_disk_tensor.py
+1
View File
@@ -164,6 +164,7 @@ class TestSafetensors(unittest.TestCase):
def test_save_all_dtypes(self):
for dtype in dtypes.fields().values():
if dtype in [dtypes.bfloat16]: continue # not supported in numpy
if dtype in [dtypes.double] and Device.DEFAULT == "METAL": continue # not supported on METAL
path = temp(f"ones.{dtype}.safetensors")
ones = Tensor(np.random.rand(10,10), dtype=dtype)
safe_save(get_state_dict(ones), path)
test/unit/test_tensor_uop_representation.py
+4 -2
View File
@@ -53,7 +53,8 @@ class TestTensorUopRepresentation(unittest.TestCase):
b = Tensor([4.,5,6]).realize()
c = a+b
print(c.lazydata)
is_pattern(c, UPat(Ops.ADD, src=(UPat(Ops.VIEW, src=(realized_pattern,)), UPat(Ops.VIEW, src=(realized_pattern,)))))
is_pattern(c, UPat(Ops.ADD, src=(realized_pattern, realized_pattern)))
#is_pattern(c, UPat(Ops.ADD, src=(UPat(Ops.VIEW, src=(realized_pattern,)), UPat(Ops.VIEW, src=(realized_pattern,)))))
def test_const_pattern(self):
a = Tensor(1)
@@ -111,7 +112,8 @@ class TestTensorUopRepresentation(unittest.TestCase):
c = a.to("TEST") # NOTE: this isn't checked
print(c.lazydata)
# TODO: COPY on a Tensor becomes a VIEW(COPY), this should be done in the scheduler not in ops
is_pattern(c, UPat(Ops.VIEW, src=(UPat(Ops.COPY, src=(UPat(Ops.DEVICE), realized_pattern,)),)))
is_pattern(c, UPat(Ops.COPY, src=(UPat(Ops.DEVICE), realized_pattern,)))
#is_pattern(c, UPat(Ops.VIEW, src=(UPat(Ops.COPY, src=(UPat(Ops.DEVICE), realized_pattern,)),)))
if __name__ == '__main__':
unittest.main()
tinygrad/ops.py
+8 -1
View File
@@ -430,7 +430,14 @@ class UOp(MathTrait, metaclass=UOpMetaClass):
# if it's a shrink, do the shrink before the copy with CONTIGUOUS
if prod(self.shape) < prod(self.base.shape): return self.contiguous().copy_to_device(device)
# COPY is COPY(DEVICE, copyin.base) -> VIEW(copyin.st)
return UOp(Ops.COPY, self.base.dtype, (UOp(Ops.DEVICE, arg=device), self.base), clone).view(unwrap(self.st))
ret = UOp(Ops.COPY, self.base.dtype, (UOp(Ops.DEVICE, arg=device), self.base), clone)
op_arg = []
mop = self
while mop is not self.base:
op_arg.append((mop.op, mop.arg))
mop = mop.src[0]
for op,arg in reversed(op_arg): ret = UOp(op, ret.dtype, (ret,), arg)
return ret
def clone(self) -> UOp: return self.copy_to_device(self.device, clone=True)
@property
def lbs(self): return [self]