Writing an autograd.Function which uses Pytorch and custom backward functionality simultaneously


I’m trying to implement a new Pytorch autograd function which should use a custom-derived gradient and combine it with a gradient obtained by the autograd engine. In essence I would like to do the following:

class some_funny_function(torch.autograd.Function):

     def forward(ctx, input):
          output1 = some_custom_function(input)
          input = torch.tensor(input, requires_grad=True)
          output2 = some_pytorch_function(input)
          ctx.save_for_backward(input, output1, output2)
          return output1, output2

     def backward(ctx, grad_input1, grad_input2):
          input, output1, output2 = ctx.saved_tensors()
          grad_output1 = custom_gradient(output1, grad_input1)
          grad_output2 = input.grad
          return grad_output1 + grad_output2

I have some input that is processed by two different functions, one I can directly write with PyTorch functions. The other can’t be done with PyTorch because it involves calling third-party software. Fortunately, I have an analytic expression for this gradient which I would like to seamlessly integrate into a PyTorch module.

Since the computation that I abbreviated ba some_pytorch_function is rather lengthy and the gradient complicated to derive, I would much rather use the autograd engine to do the hard work for me.
Unfortunately, ctx seems not to be able to do the trick for me since requires_grad seems to be disabled in a custom autograd.Function.

Do you have any suggestions on how to combine autograd and manual gradients in a nice way? Any help would be much appreciated!