From f396df26ea85d2574dad4a0a3bc6106fe71bfab4 Mon Sep 17 00:00:00 2001 From: chenyu Date: Fri, 31 Oct 2025 19:25:56 -0400 Subject: [PATCH] test custom sum (#13039) * test custom sum this is higher level than set and after? * only float --- test/test_custom_kernel.py | 14 +++++++++++++- 1 file changed, 13 insertions(+), 1 deletion(-) diff --git a/test/test_custom_kernel.py b/test/test_custom_kernel.py index ab0328b5f1..cabfcca1a8 100644 --- a/test/test_custom_kernel.py +++ b/test/test_custom_kernel.py @@ -1,6 +1,6 @@ import unittest from tinygrad import Tensor, UOp, Context -from tinygrad.uop.ops import KernelInfo, AxisType +from tinygrad.uop.ops import KernelInfo, AxisType, Ops # **** kernels **** @@ -32,6 +32,11 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp: prog = C.end(i, j) return prog.sink(arg=KernelInfo(name=f"custom_gemm_{C.shape[0]}_{C.shape[1]}_{A.shape[1]}", opts_to_apply=())) +def custom_sum(B:UOp, A:UOp) -> UOp: + # TODO: write with set and after? + i = UOp.range(A.shape[0], 0, axis_type=AxisType.REDUCE) + return B[0].store(A[i].reduce(i, arg=Ops.ADD)).sink(arg=KernelInfo(name=f"custom_sum_{A.shape[0]}", opts_to_apply=())) + # **** backward callbacks **** def backward_gemm(gradient:UOp, kernel:UOp) -> tuple[UOp, UOp]: @@ -84,6 +89,13 @@ class TestCustomKernel(unittest.TestCase): b_p1 = Tensor.custom_kernel(tst, b, fxn=custom_add_one_kernel)[0] self.assertTrue((b_p1 == 3).all().item()) + def test_sum(self): + # TODO: this only works for float, and silently fails with int + a = Tensor([1.0, 2, 3, 4, 5]) + tst = Tensor.empty(1) + b = Tensor.custom_kernel(tst, a, fxn=custom_sum)[0] + self.assertEqual(b.item(), 15) + def test_gemm(self): N = 16 a = Tensor.randn(N, N)