Function TEST

imperative/src/test/backward_graph.cpp:81–146

}

TEST(TestImperative, BackwardGraphBasic) {
    HostTensorGenerator<> gen;
    SmallVector<HostTensorND> hvs;
    SmallVector<TensorPtr> inputs;
    for (size_t i = 0; i < 2; ++i) {
        hvs.push_back(*gen({42}));
        inputs.push_back(Tensor::make(hvs.back()));
    }

    using Param = opr::Elemwise::Param;
    Param param{Param::Mode::MUL};
    auto attr = OprAttr::make("Elemwise");
    attr->cast_final_safe<OprAttr>().param.write_pod(param);

    SmallVector<LogicalTensorDesc> input_descs;
    for (auto&& i : inputs) {
        input_descs.push_back({i->layout(), i->comp_node()});
    }
    auto result = OpDef::make_backward_graph(*attr, input_descs, {true, true}, {true});
    auto&& save_for_backward = result.input_mask;
    auto&& input_has_grad = result.output_mask;

    for (size_t i = 0; i < inputs.size(); i++) {
        input_descs[i].value = inputs[i]->dev_tensor();
    }
    auto [output_descs, validated] =
            OpDef::infer_output_attrs_fallible(*attr, input_descs);
    auto outputs =
            OpDef::apply_on_physical_tensor(*attr, inputs, output_descs, validated);
    inputs.push_back(outputs[0]);
    hvs.push_back(*gen({42}));
    inputs.push_back(Tensor::make(hvs.back()));
    mgb_assert(save_for_backward.size() == inputs.size());
    for (size_t i = 0; i < inputs.size(); ++i) {
        if (!save_for_backward[i]) {
            inputs[i].reset();  // drop unused tensor
        }
    }
    SmallVector<TensorPtr> backward_graph_inputs;
    for (auto&& i : inputs) {
        if (i) {
            backward_graph_inputs.push_back(i);
        }
    }
    inputs.clear();
    auto input_grads = result.graph.apply<TensorPtr>(
            backward_graph_inputs, apply_shared_on_physical_tensor,
            [&](auto&& x) { return x; });
    mgb_assert(input_grads.size() == input_has_grad.size());
    for (size_t i = 0; i < input_has_grad.size(); ++i) {
        mgb_assert(input_has_grad[i] == static_cast<bool>(input_grads[i]));
    }

    SmallVector<HostTensorND> res;
    for (auto&& i : input_grads) {
        res.emplace_back();
        res.back().copy_from(i->dev_tensor()).sync();
    }