Function get_expected_3

tests/python/relax/test_transform_gradient_te_register.py:286–355

()

def get_expected_3():
    # fmt: off
    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def f_mul(var_A: T.handle, var_B: T.handle, var_f_mul: T.handle):
            T.func_attr({"tirx.noalias": True})
            n = T.int64()
            A = T.match_buffer(var_A, (n, n))
            B = T.match_buffer(var_B, (n, n))
            f_mul_1 = T.match_buffer(var_f_mul, (n, n))
            # with T.sblock("root"):
            for i0, i1 in T.grid(n, n):
                with T.sblock("f_mul"):
                    v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
                    T.reads(A[v_i0, v_i1], B[v_i0, v_i1])
                    T.writes(f_mul_1[v_i0, v_i1])
                    f_mul_1[v_i0, v_i1] = A[v_i0, v_i1] * B[v_i0, v_i1]

        @T.prim_func(private=True, s_tir=True)
        def f_mul_grad(var_A: T.handle, var_B: T.handle, var_C: T.handle, var_f_mul_grad_1: T.handle, var_f_mul_grad_2: T.handle):
            T.func_attr({"tirx.noalias": True})
            n = T.int64()
            A = T.match_buffer(var_A, (n, n))
            B = T.match_buffer(var_B, (n, n))
            C = T.match_buffer(var_C, (n, n))
            f_mul_grad_1 = T.match_buffer(var_f_mul_grad_1, (n, n))
            f_mul_grad_2 = T.match_buffer(var_f_mul_grad_2, (n, n))
            # with T.sblock("root"):
            for i0, i1 in T.grid(n, n):
                with T.sblock("f_mul_grad_1"):
                    v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
                    T.reads(C[v_i0, v_i1], A[v_i0, v_i1])
                    T.writes(f_mul_grad_1[v_i0, v_i1])
                    f_mul_grad_1[v_i0, v_i1] = C[v_i0, v_i1] * A[v_i0, v_i1]
            for i0, i1 in T.grid(n, n):
                with T.sblock("f_mul_grad_2"):
                    v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
                    T.reads(B[v_i0, v_i1], A[v_i0, v_i1])
                    T.writes(f_mul_grad_2[v_i0, v_i1])
                    f_mul_grad_2[v_i0, v_i1] = B[v_i0, v_i1] * A[v_i0, v_i1]

        @R.function
        def main_adjoint(a: R.Tensor(("n", "n"), dtype="float32"), b: R.Tensor(("n", "n"), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tuple(R.Tensor(("n", "n"), dtype="float32"), R.Tensor(("n", "n"), dtype="float32"))):
            n = T.int64()
            cls = Expected
            with R.dataflow():
                lv = R.call_tir(cls.f_mul, (a, b), out_sinfo=R.Tensor((n, n), dtype="float32"))
                gv: R.Tensor((), dtype="float32") = R.sum(lv, axis=None, keepdims=False)
                gv_adjoint: R.Tensor((), dtype="float32") = R.ones(R.shape([]), dtype="float32")
                lv_adjoint: R.Tensor((n, n), dtype="float32") = R.broadcast_to(gv_adjoint, R.shape([n, n]))
                lv_1 = R.call_tir(cls.f_mul_grad, (lv_adjoint, a, b), out_sinfo=[R.Tensor((n, n), dtype="float32"), R.Tensor((n, n), dtype="float32")])
                a_adjoint: R.Tensor((n, n), dtype="float32") = lv_1[0]
                b_adjoint: R.Tensor((n, n), dtype="float32") = lv_1[1]
                a_adjoint_out: R.Tensor((n, n), dtype="float32") = a_adjoint
                b_adjoint_out: R.Tensor((n, n), dtype="float32") = b_adjoint
                R.output(gv, a_adjoint_out, b_adjoint_out)
            return (gv, (a_adjoint_out, b_adjoint_out))