Method addmul_kernel

test/unit/test_gradient.py:109–114

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

class TestMultiOutputGradient(unittest.TestCase):
  @staticmethod
  def addmul_kernel(C:UOp, D:UOp, A:UOp, B:UOp) -> UOp:
    C, D, A, B = C.flatten(), D.flatten(), A.flatten(), B.flatten()
    i = UOp.range(C.numel(), 0)
    store_c = C[i].store(A[i] + B[i])
    store_d = D[i].store(A[i] * B[i])
    return UOp.group(store_c, store_d).end(i).sink(arg=KernelInfo(name="addmul")).simplify()
  @staticmethod
  def backward_addmul(grad_c, grad_d, call):
    _c, _d, a, b = call.src[1:]