Method backward

imperative/src/impl/transformations/grad.cpp:206–263

}

void GradKey::backward() {
    mgb_assert(m_frozen);
    auto& tape = m_frozen_tape;
    for (std::ptrdiff_t k = tape.size() - 1; k >= 0; --k) {
        auto& [grad_fn, op] = tape[k];
        std::string scope_name = op ? op->make_name() + ".Backward" : "CustomBackward";
        Transformation::push_scope(scope_name);
        auto grad_receiver = [&, grad_fn = grad_fn](size_t i, ValueRef grad) {
            auto& dest = grad_fn->m_dests[i];
            if (dest) {
                auto& existing_grad = dest->m_grad;
                if (!existing_grad) {
                    existing_grad = grad;
                } else {
                    existing_grad = imperative::apply(
                            ApplyOp(*Elemwise::make(Elemwise::Mode::ADD)),
                            existing_grad, grad)[0];
                }
            }
        };
        // clang-format off
        std::visit([&, grad_fn = grad_fn, op = op](auto&& backward) {
            using T = std::decay_t<decltype(backward)>;
            if constexpr (std::is_same_v<T, std::monostate>) {
                mgb_throw(AssertionError, "invalid backward");
            } else {
                // mgb_assert(grad_fn->m_slots.size() > 0);
                SmallVector<ValueRef> grads (grad_fn->m_slots.size());
                auto iter = grads.begin();
                for (auto&& slot : grad_fn->m_slots) {
                    *iter++ = slot.m_grad;
                }
                if (Profiler::is_profiling()) {
                imperative::apply(PushScope(scope_name, ScopeType::BACKWARD), Span<ValueRef>(nullptr, nullptr));
                }
                backward(grads, grad_receiver);
                if (Profiler::is_profiling()) {
                imperative::apply(PopScope(scope_name, ScopeType::BACKWARD), Span<ValueRef>(nullptr, nullptr));
                }
            }
        }, grad_fn->m_backward);
        // clang-format on
        for (auto&& dest : grad_fn->m_dests) {
            if (!dest) {
                continue;
            }
            if (!dest.m_producer_record.next && dest->callback) {
                // I'm the last grad producer, invoke callback
                if (dest->m_grad) {
                    dest->callback(dest->m_grad);
                }
            }
        }
        grad_fn->clear();
        Transformation::pop_scope(scope_name);
    }
    tape.clear();
}