| 64 | } |
| 65 | |
| 66 | pybind11::function GradKeyWrapper::get_backward_closure( |
| 67 | GradKeyWrapper* self, py::list tensors) { |
| 68 | std::vector<ValueRef> args; |
| 69 | for (auto&& tensor : tensors) { |
| 70 | auto* tw = TensorWrapper::try_cast(tensor.ptr()); |
| 71 | mgb_assert(tw, "expect Tensor"); |
| 72 | args.push_back(tw->m_tensor->data()); |
| 73 | } |
| 74 | auto closure_value = imperative::apply(GetBackwardColsure(self->m_key), args)[0]; |
| 75 | auto closure = closure_value.as_ref<FunctionValue>(); |
| 76 | auto py_function = [closure](std::vector<TensorWrapper*> tensors) { |
| 77 | std::vector<ValueRef> args; |
| 78 | for (auto* tw : tensors) { |
| 79 | args.push_back(tw->m_tensor->data()); |
| 80 | } |
| 81 | (*closure)(args); |
| 82 | }; |
| 83 | return pybind11::cpp_function(py_function); |
| 84 | } |
| 85 | |
| 86 | PyObject* GradKeyWrapper::get_name() { |
| 87 | return py::cast(m_name).release().ptr(); |