Function test_basic_fwd_with_grad

src/opr/test/loop/basic.cpp:50–107

namespace {

void test_basic_fwd_with_grad(bool dyn) {
    HostTensorGenerator<> gen;

    auto host_x = gen({23});
    auto host_loop_time =
            std::make_shared<HostTensorND>(host_x->comp_node(), dtype::Int32());
    auto graph = ComputingGraph::make();
    auto x = opr::Host2DeviceCopy::make(*graph, host_x);

    auto d = [dyn](SymbolVar var) -> SymbolVar {
        if (dyn)
            var = opr::MarkDynamicVar::make(var).node();
        return var;
    };

    auto desc_maker = [&, loop_time = opr::Host2DeviceCopy::make(
                                  *graph, host_loop_time)](LoopDesc& loop_desc) {
        auto xl = loop_desc.add_input_assignable(x).rename("xl"),
             x0 = d(loop_desc.add_input(x)).rename("x0"),
             xu = (opr::pow(d(xl), xl.make_scalar(0.9f)) * x0).rename("xu");
        loop_desc.assign(xl, xu);
        loop_desc.add_output(xu, OutputMode::LAST);
        auto cnt = d(loop_desc.get_counter_var());
        loop_desc.set_loop_condition(cnt < loop_desc.add_input(loop_time) - 1);
    };

    auto y = opr::Loop::make(desc_maker, 3)[0],
         loss = opr::reduce_sum(y, y.make_scalar(1)), gx = cg::grad(loss, x);
    HostTensorND host_y, host_gx;
    auto func = graph->compile(
            {make_callback_copy(y, host_y), make_callback_copy(gx, host_gx)});

    int& loop_time = host_loop_time->resize({1}).ptr<int>()[0];
    for (size_t sz : {12, 24}) {
        *host_x = *gen({sz});
        auto px = host_x->ptr<float>();
        for (size_t i = 0; i < sz; ++i)
            px[i] = std::abs(px[i]);

        for (loop_time = 1; loop_time <= 50; ++loop_time) {
            func->execute();
            ASSERT_EQ(host_x->shape(), host_y.shape());
            ASSERT_EQ(host_x->shape(), host_gx.shape());

            auto py = host_y.ptr<float>(), pgx = host_gx.ptr<float>();

            float dpow = 1;
            for (int i = 0; i < loop_time; ++i)
                dpow = dpow * 0.9f + 1.f;
            for (size_t i = 0; i < sz; ++i) {
                auto x = px[i], y = py[i], gx = pgx[i];
                MGB_ASSERT_FLOAT_NEAR(std::pow(x, dpow), y, 1e-5) << loop_time;
                MGB_ASSERT_FLOAT_NEAR(dpow * std::pow(x, (dpow - 1.f)), gx, 1e-5)
                        << loop_time;
            }
        }
    }
}