pytorch - 使用 pytorch 内置导数作为自定义 autograd 函数的一部分

Question

我正在寻找实现自定义 autograd。向后传递的函数是自定义函数和torch应该能够自行找到的函数的导数的混合。

举个简单的例子，假设我想为 y = x * exp(x) 创建一个函数

def custom_function(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
       ctx.save_for_backward(x)
       return x * exp(x)

    @staticmethod
    def backward(ctx, grad_output):
       x = ctx.saved_tensors[0]

       custom_derivative = x * [d/dx torch.exp(x)] + torch.exp(x)
       return grad_output * custom_derivative

如何在反向传递中调用已知 pytorch 函数的导数？