This repository was archived by the owner on Nov 17, 2023. It is now read-only.
This repository was archived by the owner on Nov 17, 2023. It is now read-only.
Second order gradient wrt inputs, expected behaviour. #14991
Description
What would be the expected behaviour of this code?
It tries to calculate the gradient of a function using the gradient wrt of the inputs of the first gradient.
def test_ag_grad():
x = mx.nd.ones((3,3))
y = mx.nd.ones((3,3))
x.attach_grad()
y.attach_grad()
with mx.autograd.record():
z = x + y
x_grad_y_grad = mx.autograd.grad(z, [x,y], create_graph=True, retain_graph=True)
print(x_grad_y_grad)
first_grad = nd.concat(*[x.reshape(-1) for x in x_grad_y_grad], dim=0)
fg_f = 2 * first_grad
second_grad = mx.autograd.grad(fg_f, [x,y], retain_graph=True)