Method bind_grad_var

src/opr/impl/loop/grad.cpp:640–683

}

bool LoopGrad::GradDesc::bind_grad_var(VarNode* owner_wrt, VarNode* owner_dest) {
    auto&& input_ogvar2info = m_fwd_graph_modifier.input_ogvar2info();
    auto info_iter = input_ogvar2info.find(owner_wrt);
    if (info_iter == input_ogvar2info.end()) {
        // caused by input vars not needed by grad
        return false;
    }
    auto&& info = m_fwd_graph_modifier.input_ogvar2info().at(owner_wrt);
    auto&& assignee2info = m_fwd_graph_modifier.assignee2info();
    bool nonzero = false;
    for (auto i : info.subgraph_var) {
        auto grad = cg::grad(m_grad_virtual_loss, i, false, false);
        if (!grad.node())
            continue;
        nonzero = true;
        bool sum_last;
        if (!assignee2info.count(i)) {
            // sum all intermediate grads
            mgb_assert(i->owner_opr()->same_type<InputMaker>());
            sum_last = false;
        } else {
            mgb_assert(!i->owner_opr()->same_type<InputMaker>());
            grad.node()->add_flag(VarNode::Flag::NO_MEM_RECLAIM);
            sum_last = true;
        }
        auto grad_sum_recorder = std::make_unique<OutputRecorderSumIntoDest>(
                sum_last, &info.grad_dest_summed, owner_dest);
        grad = grad_sum_recorder->optimize_grad_var(grad);
        do_add_output(grad, std::move(grad_sum_recorder));
        mgb_assert(m_output_record_spec_no_dedup.back()->var_sub() == grad.node());
        const_cast<OutputRecordSpecItem&>(*m_output_record_spec_no_dedup.back())
                .bind(owner_dest);
    }
    if (nonzero) {
        auto&& vec = m_uninitialized_assignor_grad_oprs;
        while (!vec.empty()) {
            auto opr = vec.back();
            vec.pop_back();
            opr->init_assignee_info(
                    m_assignor2info.at(opr->assignor()).assignees, m_grad_virtual_loss);
        }
    }
    return nonzero;
}

void LoopGrad::GradDesc::init_assignments() {
    for (auto&& i : m_fwd_graph_modifier.assignee2info()) {