Method my_kernel

test/unit/test_function.py:420–425

(C:UOp, D:UOp, A:UOp)

  def test_custom_kernel_save_unused_output(self):
    def my_kernel(C:UOp, D:UOp, A:UOp) -> UOp:
      i = UOp.range(A.shape[0], 0)
      j = UOp.range(D.shape[0], 1)
      store_c = C[i].store(A[i] * 2.0).end(i)
      store_d = D[j].store(A[j]).end(j)
      return UOp.sink(store_c, store_d, arg=KernelInfo(name="my_kernel"))

    def my_grad(d_c:UOp, call:UOp):
      a_input = call.src[3]