mirrors/tinygrad
2
0
Fork 0
mirror of https://github.com/tinygrad/tinygrad.git synced 2026-06-24 02:14:17 +00:00
Code Issues Projects Releases Packages Wiki Activity

update test_uops_stats for setitem (#14710)

Browse source 
realize both full tensor and the slice should not add to global_mem
This commit is contained in:
chenyu 2026-02-12 12:26:13 -05:00 • committed by GitHub
commit 86352988d8
No known key found for this signature in database
GPG key ID: B5690EEEBB952194
1 changed files with 3 additions and 0 deletions
Download patch file Download diff file Expand all files Collapse all files

3
test/null/test_uops_stats.py
View file

@ -72,6 +72,7 @@ class TestMemoryCount(unittest.TestCase):
t = Tensor.empty(100, dtype=dtypes.int).realize()
GlobalCounters.reset()
t[20:50] = 3
t.realize()
self.assertEqual(GlobalCounters.global_mem, 30*4) # 30 elements written
def test_setitem_slice_tensor(self):
@ -79,12 +80,14 @@ class TestMemoryCount(unittest.TestCase):
v = Tensor.empty(30, dtype=dtypes.int).realize()
GlobalCounters.reset()
t[20:50] = v
t.realize()
self.assertEqual(GlobalCounters.global_mem, 30*4*2) # 30 read + 30 written
def test_setitem_full(self):
t = Tensor.empty(100, dtype=dtypes.int).realize()
GlobalCounters.reset()
t[:] = 3
t.realize()
self.assertEqual(GlobalCounters.global_mem, 100*4) # full buffer written
@unittest.skipIf(Device.DEFAULT == "CPU", "test copy to CPU from other device")