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

Merge branch 'master' into outerworld_work

Browse Source
This commit is contained in:
George Hotz
2025-10-10 18:25:08 +08:00
committed by GitHub
parent 67a409343d ac96d98745
commit 33e8babdd8
218 changed files with 27475 additions and 295 deletions
Download Patch File Download Diff File Expand all files Collapse all files

2
test/external/process_replay/process_replay.py vendored
View File

@@ -114,7 +114,7 @@ def _pmap(fxns:dict[str, Callable]) -> None:
with multiprocessing.get_context("spawn").Pool(multiprocessing.cpu_count()) as pool:
bar = tqdm(total=row_count)
for _ in pool.imap_unordered(functools.partial(diff, fxns=fxns), range(0, row_count, PAGE_SIZE)): bar.update(PAGE_SIZE)
for _ in pool.imap_unordered(functools.partial(diff, fxns=fxns), range(0, row_count, s:=min(PAGE_SIZE, row_count))): bar.update(s)
pool.close()
pool.join()
pool.terminate()

3
test/unit/test_helpers.py
View File

@@ -1,9 +1,8 @@
import ctypes, gzip, unittest, timeit
from tinygrad import Variable
from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, CI, mv_address
from tinygrad.helpers import Context, ContextVar, argfix, colored, word_wrap, is_numpy_ndarray, CI, mv_address, get_contraction
from tinygrad.helpers import merge_dicts, strip_parens, prod, round_up, fetch, fully_flatten, from_mv, to_mv, polyN, time_to_str, cdiv, cmod, getbits
from tinygrad.tensor import Tensor, get_shape
from tinygrad.shape.view import get_contraction
import numpy as np
VARIABLE = ContextVar("VARIABLE", 0)

15
test/unit/test_shapetracker.py
View File

@@ -154,21 +154,6 @@ class TestRealStrides(unittest.TestCase):
))
self.assertEqual(st.is_expanded(), (False, False, False, True, False))
class TestRealSimplifies(unittest.TestCase):
def tearDown(self):
self.st = self.st.simplify()
assert len(self.st.views) == 1
def test_1(self):
self.st = ShapeTracker((
View.create((1, 3, 2, 11, 4, 28), (0, 308, 0, 28, 0, 1), 0, None),
View.create((1, 3, 2, 11, 26, 1, 1, 3), (0, 2464, 0, 112, 1, 0, 0, 29), 0, None)))
def test_2(self):
self.st = ShapeTracker((
View.create((8, 3, 3, 11, 2, 28), (924, 308, 0, 28, 0, 1), 0, None),
View.create((8, 1, 6, 10, 28, 3, 2, 1), (5544, 0, 0, 56, 1, 1848, 672, 0), 0, None)))
class TestIndexExpressions2d(unittest.TestCase):
def setUp(self):
shapes = [(30, 5), (15, 10), (15, 1), (5, 10), (5, 1)] # Make sure dim0 is a multiple of 5, one of the tests divides this dimension by 5

1
test/unit/test_shapetracker_math.py
View File

@@ -62,6 +62,7 @@ class TestShapeTrackerAdd(unittest.TestCase):
b = ShapeTracker.from_shape((100,))
assert a+b == b
@unittest.skip("no longer simplifies")
def test_simple_add_permute(self):
a = ShapeTracker.from_shape((10, 10))
a = a.permute((1,0))

2
test/unit/test_symbolic_shapetracker.py
View File

@@ -28,7 +28,7 @@ class TestSymbolic(unittest.TestCase):
def test_merge_view_recursion_err(self):
vm2 = View(shape=(Variable('j', 1, 10),), strides=(0,), offset=0, mask=None, contiguous=False)
vm1 = View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True)
self.assertEqual(vm2+vm1, vm1)
self.assertEqual(vm2+vm1, None)
def test_merge_view_recursion_err2(self):
vm2 = View(shape=(Variable('a', 1, 10).bind(4),), strides=(0,), offset=0, mask=None, contiguous=False)

156
test/unit/test_view.py
View File

@@ -69,161 +69,5 @@ class TestMergeDims(unittest.TestCase):
# print(f"{ShapeTracker.from_shape((2, 1, 1)).pad(((0, 0), (0, 1), (0, 1))).views[-1]}")
self.assertEqual(merge_dims((2, 2, 2), (1, 0, 0), ((0, 2), (0, 2), (0, 1))), ((2, 1, 2), (4, 0, 4)))
class TestMergeViews(unittest.TestCase):
def test_with_mask_0(self):
# from test/test_ops.py::TestOps::test_pad_reflect_mode
v0 = View(shape=(1, 1, 5, 8), strides=(0, 0, 5, 1), offset=-3, mask=((0, 1), (0, 1), (0, 5), (3, 8)), contiguous=False)
v1 = View(shape=(1, 1, 2, 2), strides=(0, 0, 8, 1), offset=3, mask=None, contiguous=False)
v = v0 + v1
self.assertIsNotNone(v)
self.assertEqual(v, View(shape=(1, 1, 2, 2), strides=(0, 0, 5, 1), offset=0, mask=None, contiguous=False))
def test_with_mask_1(self):
# from test/test_ops.py::TestOps::test_pad_reflect_mode
v0 = View(shape=(3, 3, 5, 3), strides=(27, 9, 3, 1), offset=-6, mask=((0, 3), (0, 3), (2, 4), (1, 3)), contiguous=False)
v1 = View(shape=(3, 3, 2, 2), strides=(45, 15, 3, 1), offset=7, mask=None, contiguous=False)
v = v0 + v1
self.assertIsNotNone(v)
self.assertEqual(v, View(shape=(3, 3, 2, 2), strides=(27, 9, 3, 1), offset=1, mask=None, contiguous=False))
def test_with_mask_2(self):
# from test/test_ops.py::TestOps::test_pad_reflect_mode
v0 = View(shape=(3, 3, 5, 3), strides=(27, 9, -3, 1), offset=6, mask=((0, 3), (0, 3), (0, 2), (0, 2)), contiguous=False)
v1 = View(shape=(3, 3, 2, 2), strides=(45, 15, -3, 1), offset=3, mask=None, contiguous=False)
v = v0 + v1
self.assertIsNotNone(v)
self.assertEqual(v, View(shape=(3, 3, 2, 2), strides=(27, 9, 3, 1), offset=3, mask=None, contiguous=False))
def test_with_mask_3(self):
# from test/test_ops.py::TestOps::test_pad_reflect_mode
# has a mask in the final view
v0 = View(shape=(3, 3, 4, 4), strides=(27, 9, 3, 1), offset=-5, mask=((0, 3), (0, 3), (2, 4), (0, 2)), contiguous=False)
v1 = View(shape=(3, 3, 4, 2), strides=(48, 16, 4, 1), offset=0, mask=None, contiguous=False)
v = v0 + v1
self.assertIsNotNone(v)
self.assertEqual(v, View(shape=(3, 3, 4, 2), strides=(27, 9, 3, 1), offset=-5, mask=((0, 3), (0, 3), (2, 4), (0, 2)), contiguous=False))
def test_with_mask_4(self):
# from test/test_ops.py::TestOps::test_pad_reflect_mode
# has a mask in the final view
v0 = View(shape=(3, 3, 5, 3), strides=(27, 9, -3, 1), offset=6, mask=((0, 3), (0, 3), (0, 2), (1, 3)), contiguous=False)
v1 = View(shape=(3, 3, 3, 3), strides=(45, 15, 3, 1), offset=6, mask=None, contiguous=False)
v = v0 + v1
self.assertIsNotNone(v)
self.assertEqual(v, View(shape=(3, 3, 3, 3), strides=(0, 0, 0, 0), offset=0, mask=((0, 0), (0, 0), (0, 0), (0, 0)), contiguous=False))
def test_with_mask_5(self):
# from test/test_ops.py::TestOps::test_pad_reflect_mode
# has a mask in the final view
v0 = View(shape=(1, 1, 6, 5), strides=(0, 0, 5, 1), offset=-5, mask=((0, 1), (0, 1), (1, 6), (0, 5)), contiguous=False)
v1 = View(shape=(1, 1, 6, 3), strides=(0, 0, 5, -1), offset=3, mask=None, contiguous=False)
v = v0 + v1
self.assertIsNotNone(v)
self.assertEqual(v, View(shape=(1, 1, 6, 3), strides=(0, 0, 5, -1), offset=-2, mask=((0, 1), (0, 1), (1, 6), (0, 3)), contiguous=False))
@unittest.expectedFailure # TODO: fix these
def test_merges_from_fuzzer1(self):
v0 = View(shape=(2, 4), strides=(2, 1), offset=-2, mask=((0, 2), (2, 4)), contiguous=False)
v1 = View(shape=(2, 4, 2, 2), strides=(4, 0, -2, -1), offset=3, mask=None, contiguous=False)
target = View(shape=(2, 4, 2, 2), strides=(2, 0, 0, -1), offset=1, mask=((0, 2), (0, 4), (0, 1), (0, 2)), contiguous=False)
v = v0 + v1
self.assertIsNotNone(v)
self.assertEqual(v, target)
@unittest.expectedFailure # TODO: fix these
def test_merges_from_fuzzer2(self):
v0 = View(shape=(5, 10, 12), strides=(100, 1, 10), offset=-20, mask=((0, 5), (0, 10), (2, 12)), contiguous=False)
v1 = View(shape=(10, 6, 5, 2, 2), strides=(12, 2, 120, 1, 0), offset=0, mask=None, contiguous=False)
target = View(shape=(10, 6, 5, 2, 2), strides=(1, 20, 100, 10, 0), offset=-20, mask=((0, 10), (1, 6), (0, 5), (0, 2), (0, 2)), contiguous=False)
v = v0 + v1
self.assertIsNotNone(v)
self.assertEqual(v, target)
@unittest.expectedFailure # TODO: fix these
def test_merges_from_fuzzer3(self):
v0 = View(shape=(8, 7, 3), strides=(1, 12, -4), offset=6, mask=((2, 6), (0, 7), (0, 3)), contiguous=False)
v1 = View(shape=(4, 2, 6, 2, 1), strides=(42, 21, 3, 1, 0), offset=4, mask=None, contiguous=False)
target = View(shape=(4, 2, 6, 2, 1), strides=(2, 1, 12, -4, 0), offset=14, mask=((1, 3), (0, 2), (0, 6), (0, 2), (0, 1)), contiguous=False)
v = v0 + v1
self.assertIsNotNone(v)
self.assertEqual(v, target)
@unittest.expectedFailure # TODO: fix these
def test_merges_from_fuzzer4(self):
v0 = View(shape=(7, 21, 3), strides=(54, 3, 1), offset=-9, mask=((0, 6), (3, 21), (0, 3)), contiguous=False)
v1 = View(shape=(5, 3, 3, 7), strides=(63, 1, 3, 9), offset=63, mask=None, contiguous=False)
target = View(shape=(5, 3, 3, 7), strides=(54, 1, 3, 9), offset=45, mask=((0, 5), (0, 3), (0, 3), (1, 7)), contiguous=False)
v = v0 + v1
self.assertIsNotNone(v)
self.assertEqual(v, target)
@unittest.expectedFailure # TODO: fix these
def test_merges_from_fuzzer5(self):
v0 = View(shape=(5, 1, 24), strides=(20, 0, 1), offset=-2, mask=((0, 5), (0, 1), (2, 22)), contiguous=False)
v1 = View(shape=(12, 2, 5, 2, 1), strides=(2, 1, 24, 0, 0), offset=0, mask=None, contiguous=False)
target = View(shape=(12, 2, 5, 2, 1), strides=(2, 1, 20, 0, 0), offset=-2, mask=((1, 11), (0, 2), (0, 5), (0, 2), (0, 1)), contiguous=False)
v = v0 + v1
self.assertIsNotNone(v)
self.assertEqual(v, target)
def test_merge_views_variable(self):
from tinygrad import Variable
N = 100
start_pos = Variable("start_pos", 1, N-1)
v0 = View(shape=(N, 32, 2), strides=(32, 1, 0), offset=0, mask=((0, N), (0, 32), (0, 1)), contiguous=False)
v1 = View(shape=(1, 8, 1, 32), strides=(0, 0, 0, 2), offset=start_pos*64, mask=None, contiguous=False)
target = View(shape=(1, 8, 1, 32), strides=(0,0,0,1), offset=start_pos*32, mask=None, contiguous=False)
v = v0 + v1
self.assertIsNotNone(v)
self.assertEqual(v, target)
def test_view_padded_area1(self):
# test_multinomial
v0 = View(shape=(2,), strides=(0,), offset=0, mask=((1, 2),), contiguous=False)
v1 = View(shape=(1,), strides=(0,), offset=0, mask=None, contiguous=True)
v = v0 + v1
self.assertIsNotNone(v)
self.assertEqual(v, View(shape=(1,), strides=(0,), offset=0, mask=((0, 0),), contiguous=False))
def test_view_padded_area2(self):
# test_pad_reflect_mode
v0 = View(shape=(1, 1, 10, 7), strides=(0, 0, 5, 1), offset=-15, mask=((0, 1), (0, 1), (3, 8), (0, 5)), contiguous=False)
v1 = View(shape=(0, 0, 0, 0), strides=(0, 0, 0, 0), offset=0, mask=None, contiguous=True)
v = v0 + v1
self.assertIsNotNone(v)
self.assertEqual(v, View(shape=(0, 0, 0, 0), strides=(0, 0, 0, 0), offset=0, mask=None, contiguous=True))
def test_view_padded_area3(self):
# test_roll
v0 = View(shape=(2, 4), strides=(0, 1), offset=4, mask=((0, 1), (0, 4)), contiguous=False)
v1 = View(shape=(1, 4), strides=(0, 1), offset=4, mask=None, contiguous=False)
v = v0 + v1
self.assertIsNotNone(v)
self.assertEqual(v, View(shape=(1, 4), strides=(0, 0), offset=0, mask=((0, 0), (0, 0)), contiguous=False))
def test_view_padded_area4(self):
# test_std_mean
v0 = View(shape=(2,), strides=(0,), offset=0, mask=((0, 1),), contiguous=False)
v1 = View(shape=(1, 1, 1), strides=(0, 0, 0), offset=1, mask=None, contiguous=False)
v = v0 + v1
self.assertIsNotNone(v)
self.assertEqual(v, View(shape=(1, 1, 1), strides=(0, 0, 0), offset=0, mask=((0, 0), (0, 0), (0, 0)), contiguous=False))
def test_empty_shape_view1(self):
# test_stack_slice
v0 = View(shape=(3, 5), strides=(0, 1), offset=0, mask=((0, 1), (0, 5)), contiguous=False)
v1 = View(shape=(), strides=(), offset=0, mask=None, contiguous=True)
v = v0 + v1
self.assertIsNotNone(v)
self.assertEqual(v, View(shape=(), strides=(), offset=0, mask=None, contiguous=True))
def test_empty_shape_view2(self):
# test_std_mean
v0 = View(shape=(2,), strides=(0,), offset=0, mask=((1, 2),), contiguous=False)
v1 = View(shape=(), strides=(), offset=0, mask=None, contiguous=True)
v = v0 + v1
# TODO: why is this different?
self.assertIsNone(v)
if __name__ == '__main__':
unittest.main()