From bebec73471bd98a992ede37bc9d472e281c0dc60 Mon Sep 17 00:00:00 2001 From: chenyu Date: Sat, 1 Nov 2025 10:45:30 -0400 Subject: [PATCH] write custom_sum with set and after (#13045) --- test/test_custom_kernel.py | 7 ++++--- 1 file changed, 4 insertions(+), 3 deletions(-) diff --git a/test/test_custom_kernel.py b/test/test_custom_kernel.py index a1666aba2a..b779ab3868 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, Ops +from tinygrad.uop.ops import KernelInfo, AxisType # **** kernels **** @@ -33,9 +33,10 @@ def custom_gemm(C:UOp, A:UOp, B:UOp) -> UOp: 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=())) + B = B[0].set(0.0) + B = B[0].set(B.after(i)[0] + A[i], end=i) + return B.sink(arg=KernelInfo(name=f"custom_sum_{A.shape[0]}", opts_to_apply=())) def flip_contract_kernel(dest:UOp, src:UOp): assert dest.size%4 == 0