Function test_null_grad

src/opr/test/loop/basic.cpp:1353–1494

namespace {

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

    auto zg = [](SymbolVar var) {
        return opr::SetGrad::make(var, opr::SetGrad::zero_grad);
    };

    auto spow = [&](SymbolVar a, float p) { return opr::pow(d(a), a.make_scalar(p)); };

    constexpr int LOOP_TIME = 1;
    HostTensorGenerator<dtype::Float32, RandomDistribution::UNIFORM> gen{1e-2, 1};
    constexpr size_t SIZE0 = 23, SIZE1 = 32;
    auto host_x0 = gen({SIZE0}), host_x1 = gen({SIZE1});

    auto graph = ComputingGraph::make();
    auto outgraph_x0 = opr::Host2DeviceCopy::make(*graph, host_x0),
         outgraph_x1 = opr::Host2DeviceCopy::make(*graph, host_x1);

    auto desc_maker = [&](LoopDesc& desc) {
        auto x0a = desc.add_input_assignable(outgraph_x0).rename("x0a"),
             x0b = desc.add_input_assignable(outgraph_x0).rename("x0b"),

             // one value as multiple assignors
                val0 = (spow(x0a, 0.4) * spow(x0b, 0.5)).rename("val0"),

             x0c = desc.add_input_assignable(outgraph_x0).rename("x0c"),
             x0c1 = desc.add_input(outgraph_x0).rename("x0c1"),
             // assignee null grad
                val1 = (spow(zg(x0c), 0.5) * x0c1).rename("val1"),

             x0d = desc.add_input_assignable(outgraph_x0).rename("x0d"),
             // assignor null grad
                val2 = (spow(x0d, 0.6) * x0c1).rename("val2"),

             // null outgrad in par graph
                x1 = desc.add_input(outgraph_x1).rename("x1"),
             x1a = desc.add_input_assignable(outgraph_x1).rename("x1a"),
             val3 = (spow(x1a, 0.7) * x1).rename("val3");

        desc.assign(x0a, val0);
        desc.assign(x0b, val0);
        desc.assign(x0c, val1);
        desc.assign(x0d, val2);
        desc.assign(x1a, val3);

        desc.add_output(val0, OutputMode::LAST);
        desc.add_output(val1, OutputMode::LAST);
        desc.add_output(zg(val2), OutputMode::LAST);
        desc.add_output(val3, OutputMode::LAST);

        desc.set_loop_condition(desc.get_counter_var() < LOOP_TIME - 1);
    };

    auto y = opr::Loop::make(desc_maker);