Suppose I have a bivariate covariance function c(x,y) and at input [x1,...,xn] and [y1,...,ym] and I would like to compute the covariance matrix with i,jth element d/dx c(xi,yj). Here I have assumed that each xi and yj are scalar valued, however this could easily be extended to finding the kth partial derivative of the covariance matrix. In general, I would like this to be extended to higher derivatives too (e.g. I would like to compute the matrix with i,jth element d/dx d/dy c(xi,yj)), but this may be an easy generalisation of the first issue.
I feel that PyTorch’s automatic differentiation system should be able to handle this, but I haven’t been able to come up with any nice solution. I think the issue is to do with broadcasting derivatives. For instance, to compute the non-differentiated covariance matrix, I would do something like c(x.unsqueeze(1),y),  but the following code doesn’t work as I would like (using autograd.grad):
X = torch.Tensor([1,2,3]).requires_grad_(True).unsqueeze(1)
Y = torch.Tensor([4,5]).requires_grad_(True)
cov_matrix = c(X,Y)
dx_cov_matrix = grad(cov_matrix, [X], grad_outputs=[torch.ones([X.size(0),Y.size(0)])], create_graph=True)

Any help would be appreciated! Thanks