opened 03:52PM - 30 Jul 26 UTC
closed 03:53PM - 30 Jul 26 UTC
bot-triaged
### 📚 The doc issue
I was reading https://docs.pytorch.org/docs/2.13/notes/exte… nding.html
the code is
```
class MyCube(torch.autograd.Function):
@ staticmethod
def forward(x):
# We wish to save dx for backward. In order to do so, it must
# be returned as an output.
dx = 3 * x ** 2
result = x ** 3
return result, dx
@ staticmethod
def setup_context(ctx, inputs, output):
x, = inputs
result, dx = output
ctx.save_for_backward(x, dx)
@ staticmethod
def backward(ctx, grad_output, grad_dx):
x, dx = ctx.saved_tensors
# In order for the autograd.Function to work with higher-order
# gradients, we must add the gradient contribution of `dx`,
# which is grad_dx * 6 * x.
result = grad_output * dx + grad_dx * 6 * x
return result
# Wrap MyCube in a function so that it is clearer what the output is
def my_cube(x):
result, dx = MyCube.apply(x)
return result
```
(copied from your docs web)
But when I call
```dx.backward(gradient=torch.tensor([7.]), inputs=[input])```
or
```result.backward(gradient=torch.tensor([3.]), inputs=[input])```
I got only one meaningful input param in the backward function.
my code:
```
class MyCube(torch.autograd.Function):
@staticmethod
def forward(x):
# We wish to save dx for backward. In order to do so, it must
# be returned as an output.
dx = 3 * x ** 2
result = x ** 3
return result, dx
@staticmethod
def setup_context(ctx, inputs, output):
x, = inputs
result, dx = output
ctx.save_for_backward(x, dx)
@staticmethod
def backward(ctx, grad_output, grad_dx):
x, dx = ctx.saved_tensors
# In order for the autograd.Function to work with higher-order
# gradients, we must add the gradient contribution of `dx`,
# which is grad_dx * 6 * x.
result = grad_output * dx + grad_dx * 6 * x
return result
# Wrap MyCube in a function so that it is clearer what the output is
input = torch.tensor([2.], requires_grad=True)
result:torch.Tensor
result, dx = MyCube.apply(input)
dx.backward(gradient=torch.tensor([7.]), inputs=[input])
#result.backward(gradient=torch.tensor([3.]), inputs=[input]) # this line raises exception.
```
If I call from dx, then result is [0.], if I call from result, then dx is [0.].
And dx.retail_grad() or result.retail_grad() don't do anything.
Let's ignore the math meaning. How do I get both in the backward function, or how do I call the backward twice?
Thank you.
### Suggest a potential alternative/fix
Idk. A working example would be nice.
An example would be very helpful. Thank you.
Hi,retain_graph=True during your first call to backward (otherwise the graph will be freed, to save memory, while in reality you need still need it for the second call).
So the last two lines should be:
dx.backward(gradient=torch.tensor([7.]), inputs=[input], retain_graph=True)
result.backward(gradient=torch.tensor([3.]), inputs=[input])
This should fix your exception.
This will have the effect of storing in input.grad the gradient of dx wrt input multiplied by 7 + the gradient of result wrt input multiplied by 3.
To make it more efficient, you can also compute that in a single torch.autograd.backward call:
torch.autograd.backward(tensors=[dx, result], grad_tensors=[torch.tensor([7.]), torch.tensor([3.])])
Also, you mentioned:
And dx.retail_grad() or result.retail_grad() don’t do anything.
retain_grad is a method you can call on individual tensors, that will ensure they store a gradient wrt themselves in their own .grad field despite not being leaves in the computation graph (see this ). In reality, you rarely need to use this. This is very different from the retain_graph I’m talking about, which is a parameter of the backward and grad functions, telling autograd to not free the computation graph yet, and which you have to use when you make multiple calls to backward or grad.
I hope this helps!