nimlgen/tinygrad
1
0
Fork 0
forked from tinygrad/tinygrad
Code Pull Requests Activity

test permuted input in custom Ops.PROGRAM test (#17279)

Browse Source 
* test permuted input in custom Ops.PROGRAM test

* dtype
This commit is contained in:
qazal
2026-07-29 18:29:05 +09:00
committed by GitHub
parent 527e57300c
commit 6ea7d366fa
1 changed files with 5 additions and 8 deletions
Download Patch File Download Diff File Expand all files Collapse all files
test/backend/test_custom_kernel.py
+5 -8
View File
@@ -421,19 +421,16 @@ class TestCustomKernel(unittest.TestCase):
@Context(DEV="CPU")
def test_simple_from_source(self):
a = Tensor([0., 1., 2.]).realize()
src = "void test_src(float* restrict a) { a[0] = 1.0; }"
a = Tensor.arange(4).clone().realize()
src = "void test_src(int* restrict a) { a[0] = 1; }"
# TODO: it currently requires a compiler for Ops.BINARY
from tinygrad.device import Device
binary = Device[a.device].renderer.compiler.compile(src)
def custom_src_kernel(A:UOp) -> UOp:
sink = UOp.sink(A, arg=KernelInfo(name="test_src"))
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())),
UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))
a = Tensor.custom_kernel(a, fxn=custom_src_kernel)[0]
self.assertEqual(a.tolist(), [1., 1., 2.])
return UOp(Ops.PROGRAM, src=(sink, UOp(Ops.LINEAR, src=tuple(sink.toposort())), UOp(Ops.SOURCE, arg=src), UOp(Ops.BINARY, arg=binary)))
a = Tensor.custom_kernel(a.reshape(2, 2).T, fxn=custom_src_kernel)[0]
self.assertEqual(a.tolist(), [[1, 2], [1, 3]])
class TestUOpReduce(unittest.TestCase):
def test_uop_sum(self):