I want to use the partial derivative of a function, but I don’t want to manually calculate it so I use autograd to calculate it for me.
This seems to work (?) see below.
Problem: While calculating the partial derivative w.r.t x of f(x,y) it also calculates the partial derivative w.r.t all parents of both x and y while I only want to do it for x in the forward pass. I can overcome this by zeroing out gradients but it still seems like a waste of work. How can I do this better?
Here f(x,y) = (x*y)^2
import torch
print(torch.__version__)
#> '0.4.0a0+059299b'

from torch.autograd import Variable
from torch import nn
from torch.nn.parameter import Parameter
import torch.optim as optim

def dfdx(x,param,f):
    x = Variable(x.data,requires_grad=True)
    z = f(x,param)
    z.backward([torch.ones_like(z)],create_graph=True)
    dfdx = x.grad

    return dfdx

n_batch = 5
x = Variable(3*torch.ones(n_batch), requires_grad = False)
w = Parameter(.5*torch.ones(1), requires_grad = True)
p = Variable(torch.ones(n_batch))*w
# Forward pass 
dfdx_val = dfdx(x,p,lambda x,p : (p*x)**2)
print('df/dx')
print('result :',dfdx_val.data.numpy().T) 
print('sought :',((p*p)*2*x).data.numpy().T) # (p*p)*2*x 
# > df/dx
# > result : [[1.5 1.5 1.5 1.5 1.5]]
# > sought : [[1.5 1.5 1.5 1.5 1.5]]

print('w.grad :',w.grad) # <-- PROBLEM. I didn't want this to be calculated.
# > w.grad : Variable containing:
# > 135

# Backward pass; 
# Since a .backward() is called during forward in dfdx, 
# need to call it *After* Forward
w.grad.data.zero_() 
dfdx_val.mean().backward() # will only work on scalar! 
print('df/dxdp')
print('result :',(w.grad).data.numpy().T) 
print('sought :',(4*w*x).data.numpy()[:1].T) # 4*p*x 
# > df/dxdp
# > result : [[6.]]
# > sought : [[6.]]

It might help to think of y in f(x,y) as input from layers above while x is a query.
Keywords; grad of grad, higher order derivatives, partial derivative as functional/layer…