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

assign into empty works (#15256)

Browse Source
This commit is contained in:
Sieds Lykles
2026-03-13 15:24:29 +01:00
committed by GitHub
parent 60b1b908c6
commit 4b59083d7c
2 changed files with 2 additions and 3 deletions
Download Patch File Download Diff File Expand all files Collapse all files

3
test/unit/test_setitem_schedule.py
View File

@@ -36,8 +36,7 @@ class TestSetitemInto(unittest.TestCase):
self.assertEqual(GlobalCounters.kernel_count, 0)
t.realize()
self.assertEqual(GlobalCounters.kernel_count, 1)
# TODO: this can be just 4 if empty goes through is_realized setitem path
self.assertEqual(GlobalCounters.global_mem, 4*(3*2+1)) # 3 elements had +1, 1 is assigned directly
self.assertEqual(GlobalCounters.global_mem, 4)
t[1].realize()
t.realize()
self.assertEqual(GlobalCounters.kernel_count, 1)

2
tinygrad/tensor.py
View File

@@ -1347,7 +1347,7 @@ class Tensor(OpMixin):
if is_disk: raise RuntimeError("advanced setitem is not supported for DISK tensors")
if not isinstance(v, Tensor): v = Tensor(v, device=self.device, dtype=self.dtype)
self.assign(self._getitem(indices, v))
elif is_disk or self.uop.is_realized or self.uop.base.op is Ops.AFTER: # basic setitem, self is realized
elif is_disk or self.uop.is_realized or self.uop.base.op in (Ops.AFTER, Ops.BUFFER): # basic setitem, self is realized
view = self[indices]
if isinstance(v, Tensor) and v.uop.op is Ops.ASSIGN and v.uop in view.uop.base.src: return
view.assign(v)