Method attach

imperative/python/src/grad.cpp:25–48

GradKeyWrapper::GradKeyWrapper() {}

void GradKeyWrapper::attach(PyObject* const* args, size_t nargs) {
    if (nargs != 2) {
        throw py::type_error("expect 2 arguments");
    }
    auto* tw = TensorWrapper::try_cast(args[0]);
    if (!tw) {
        throw py::type_error("argument 1 must be Tensor");
    }
    py::object callback;
    if (args[1] != Py_None) {
        callback = py::reinterpret_borrow<py::object>(args[1]);
    }
    GenericFunction generic_callback = [=](Span<ValueRef> inputs) -> ValueRefList {
        mgb_assert(inputs.size() == 1);
        if (callback) {
            callback(TensorWrapper::make(py_tensor_type, inputs[0]));
        }
        return {};
    };
    auto attached_value = imperative::apply(
            AttachGrad(m_key), tw->m_tensor->data(),
            FunctionValue::make(generic_callback))[0];
    tw->m_tensor->reset(attached_value);
}

void GradKeyWrapper::backward(GradKeyWrapper* self, py::list tensors, py::list grads) {
    std::vector<ValueRef> args;