How to get the full Jacobian of a derivative in PyTorch?
Lets consider a simple tensor x and lets define another one which depends on x and have multiple dimension : y = (x, 2x, x^2).

How can I have the full gradient dy/dx = (1,2,x) ? 

For example lets take the code :

import torch
from torch.autograd import grad

x = 2 * torch.ones(1)
x.requires_grad = True
y = torch.cat((x, 2*x, x*x))
# dy_dx = ???


This is what I have I unsuccessfuly tried so far :

&gt;&gt;&gt; dy_dx = grad(y, x, grad_outputs=torch.ones_like(y), create_graph=True)
(tensor([7.], grad_fn=&lt;AddBackward0&gt;),)
&gt;&gt;&gt; dy_dx = grad(y, x, grad_outputs=torch.Tensor([1,0,0]), create_graph=True)
(tensor([1.], grad_fn=&lt;AddBackward0&gt;),)
&gt;&gt;&gt; dy_dx = grad(y, [x,x,x], grad_outputs=torch.eye(3), create_graph=True)
(tensor([7.], grad_fn=&lt;AddBackward0&gt;),)


Each times I got only part of the gradient or an accumulated version...

I know I could use a for loop using the second expression like

dy_dx = torch.zeros_like(y)
coord = torch.zeros_like(y)
for i in range (y.size(0)):
    coord[i] = 1
    dy_dx[i], = grad(y, x, grad_outputs=coord, create_graph=True)
    coord[i] = 0


However, as I am handeling with high dimensions tensors, this for loop could take too much time to compute. Moreover, there must be a way to perform the full jacobian without acuumulating the gradient...

Does anyone has the solution ? Or an alternative ?

torch.autograd.grad in PyTorch is aggregated. To have a vector auto-differentiated with respect to the input, use torch.autograd.functional.jacobian.
