Function get_expected_2

tests/python/relax/test_transform_gradient_te_register.py:177–226

()

def get_expected_2():
    # 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"), f_mul2: 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_mul2"):
                    v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
                    T.reads(A[v_i0, v_i1])
                    T.writes(f_mul2[v_i0, v_i1])
                    f_mul2[v_i0, v_i1] = A[v_i0, v_i1] * T.float32(2)

        @T.prim_func(private=True, s_tir=True)
        def f_mulk_grad(A: T.Buffer((T.int64(5), T.int64(5)), "float32"), B: T.Buffer((T.int64(5), T.int64(5)), "float32"), f_mulk_grad_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_mulk_grad"):
                    v_i0, v_i1 = T.axis.remap("SS", [i0, i1])
                    T.reads(A[v_i0, v_i1])
                    T.writes(f_mulk_grad_1[v_i0, v_i1])
                    f_mulk_grad_1[v_i0, v_i1] = A[v_i0, v_i1] * T.float32(2)

        @R.function
        def main_adjoint(a: R.Tensor((5, 5), dtype="float32")) -> R.Tuple(R.Tensor((), dtype="float32"), R.Tuple(R.Tensor((5, 5), dtype="float32"))):
            cls = Expected
            with R.dataflow():
                lv = R.call_tir(cls.f_mul, (a,), 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_mulk_grad, (lv_adjoint, a), out_sinfo=R.Tensor((5, 5), dtype="float32"))
                a_adjoint: R.Tensor((5, 5), dtype="float32") = lv_1
                a_adjoint_out: R.Tensor((5, 5), dtype="float32") = a_adjoint
                R.output(gv, a_adjoint_out)
            return (gv, (a_adjoint_out,))

        @R.function
        def main(a: 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,), out_sinfo=R.Tensor((5, 5), dtype="float32"), te_grad_name="f_mulk_grad", te_grad_kwargs={"k": T.float32(2)})
                gv: R.Tensor((), dtype="float32") = R.sum(lv, axis=None, keepdims=False)
                R.output(gv)
            return gv
    # fmt: on
    return Expected


def test_emit_te_kwargs(register_te_grads):