#include <torch/extension.h> torch::Tensor my_add_forward(const torch::Tensor& input1, const torch::Tensor& input2) { return input1 + input2; } std::vector<torch::Tensor> my_add_backward(const torch::Tensor& grad_output) { return {grad_output, grad_output}; } PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) { m.def("forward", &my_add_forward, "MyAdd forward"); m.def("backward", &my_add_backward, "MyAdd backward"); }

这段代码是使用 PyTorch C++ 扩展实现的一个简单的张量加法操作。其中，my_add_forward 函数接受两个张量参数 input1 和 input2，返回它们的和。my_add_backward 函数接受一个张量参数 grad_output，返回一个包含两个张量的 vector，这两个张量都是 grad_output。最后，在 PYBIND11_MODULE 宏中，我们将这两个函数绑定到 Python 模块中，使得我们可以在 Python 中调用它们。