Method make_backward_closure

imperative/src/impl/transformations/grad.cpp:558–583

}

GenericFunction GradTransformation::make_backward_closure(Span<ValueRef> ys) {
    // reset GradKey
    auto grad_key = m_key;
    std::vector<GradSlotPtr> y_slots;
    for (auto&& y : ys) {
        if (auto&& grad_value = as_grad_value(y)) {
            y_slots.push_back(grad_value->slot());
        } else {
            y_slots.emplace_back();
        }
    }
    GenericFunction closure = [grad_key, y_slots](Span<ValueRef> dys) -> ValueRefList {
        size_t nr_grads = y_slots.size();
        mgb_assert(dys.size() == nr_grads);
        for (size_t i = 0; i < nr_grads; ++i) {
            if (y_slots[i]) {
                y_slots[i]->m_grad = dys[i];
            }
        }
        grad_key->backward();
        return {};
    };
    grad_key->freeze();
    cleanup();
    return closure;
}

void GradTransformation::on_unregister() noexcept {
    cleanup();