Использование Pytorch, встроенного в производную, как часть пользовательской функции автограда

Question

Я ищу реализацию собственного автограда. Функция, в которой обратный проход представляет собой смесь пользовательской функции и производной функции, факел которой должен быть в состоянии найти сам.

В качестве простого примера, скажем, я хотел создать функцию для 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 в обратном проходе?