superrob/IQ.Pilot
1
0
Fork 0
forked from IQ.Lvbs/IQ.Pilot
Code Pull Requests Activity

IQ.Pilot Release Commit @ f2a861c

Browse Source
This commit is contained in:
IQ.Lvbs CI [bot]
2026-09-02 15:07:09 -05:00
parent b42569dbca
commit e8748fd704
5497 changed files with 316070 additions and 179848 deletions
Download Patch File Download Diff File Expand all files Collapse all files

21
artifacts/package_sources/tinygrad/test/unit/test_getitem_ops.py Normal file
View File

@@ -0,0 +1,21 @@
import unittest
import numpy as np
from tinygrad import Tensor, GlobalCounters
class TestGetitemOps(unittest.TestCase):
def test_two_tensor_indices(self):
# linear indexing is O(idx_size), one-hot masks is O(idx_size * src_size)
src_np = np.random.rand(10, 100, 200).astype(np.float32)
idx1_np, idx2_np = np.random.randint(0, 100, (50, 60), dtype=np.int32), np.random.randint(0, 200, (50, 60), dtype=np.int32)
src, idx1, idx2 = Tensor(src_np), Tensor(idx1_np), Tensor(idx2_np)
# O(50*60) = 3K vs O(50*60*100*200) = 60M
GlobalCounters.reset()
np.testing.assert_allclose(src_np[0, idx1_np, idx2_np], src[0, idx1, idx2].numpy())
self.assertLess(GlobalCounters.global_ops, 100_000)
# consecutive indices not starting from dim 0: O(10*50*60) = 30K vs O(10*50*60*100*200) = 600M
GlobalCounters.reset()
np.testing.assert_allclose(src_np[:, idx1_np, idx2_np], src[:, idx1, idx2].numpy())
self.assertLess(GlobalCounters.global_ops, 1_000_000)
if __name__ == '__main__':
unittest.main()