Function get_expected_1

tests/python/relax/test_transform_gradient_te_register.py:61–118

()

def get_expected_1():
    # fmt: off
    @I.ir_module(s_tir=True)
    class Expected:
        @T.prim_func(private=True, s_tir=True)
        def f_mul(A: T.Buffer((T.int64(5), T.int64(5)), "float32"), B: T.Buffer((T.int64(5), T.int64(5)), "float32"), f_mul_1: T.Buffer((T.int64(5), T.int64(5)), "float32")):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            for i0, i1 in T.grid(T.int64(5), T.int64(5)):
                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(A: T.Buffer((T.int64(5), T.int64(5)), "float32"), B: T.Buffer((T.int64(5), T.int64(5)), "float32"), C: T.Buffer((T.int64(5), T.int64(5)), "float32"), f_mul_grad_1: T.Buffer((T.int64(5), T.int64(5)), "float32"), f_mul_grad_2: T.Buffer((T.int64(5), T.int64(5)), "float32")):
            T.func_attr({"tirx.noalias": True})
            # with T.sblock("root"):
            for i0, i1 in T.grid(T.int64(5), T.int64(5)):
                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(T.int64(5), T.int64(5)):
                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((5, 5), dtype="float32"), b: R.Tensor((5, 5), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tuple(R.Tensor((5, 5), dtype="float32"), R.Tensor((5, 5), dtype="float32"))):
            cls = Expected
            with R.dataflow():
                lv = R.call_tir(cls.f_mul, (a, b), out_sinfo=R.Tensor((5, 5), 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((5, 5), dtype="float32") = R.broadcast_to(gv_adjoint, R.shape([5, 5]))
                lv_1 = R.call_tir(cls.f_mul_grad, (lv_adjoint, a, b), out_sinfo=[R.Tensor((5, 5), dtype="float32"), R.Tensor((5, 5), dtype="float32")])
                a_adjoint: R.Tensor((5, 5), dtype="float32") = lv_1[0]
                b_adjoint: R.Tensor((5, 5), dtype="float32") = lv_1[1]
                a_adjoint_out: R.Tensor((5, 5), dtype="float32") = a_adjoint
                b_adjoint_out: R.Tensor((5, 5), dtype="float32") = b_adjoint
                R.output(gv, a_adjoint_out, b_adjoint_out)
            return (gv, (a_adjoint_out, b_adjoint_out))

        @R.function
        def main(a: R.Tensor((5, 5), dtype="float32"), b: R.Tensor((5, 5), dtype="float32")) -> R.Tensor((), dtype="float32"):
            cls = Expected
            with R.dataflow():
                lv = R.call_tir_with_grad(cls.f_mul, (a, b), out_sinfo=R.Tensor((5, 5), dtype="float32"), te_grad_name="f_mul_grad")
                gv: R.Tensor((), dtype="float32") = R.sum(lv, axis=None, keepdims=False)
                R.output(gv)
            return gv
    # fmt: on
    return Expected