Potentially silently incorrect behavior with functorch <-> autograd.Function interaction

[zou3519]

🐛 Describe the bug

Initially reported by @AlphaBetaGamma96. Repro script over at https://gist.github.com/zou3519/8b25cc439cf5c51152f0d35a06cce010 . I'm still working on making it smaller, but just putting this issue up for tracking.

Versions

2.0

cc @ezyang @gchanan @albanD @gqchen @pearu @nikitaved @soulitzer @lezcano @Varal7 @Chillee @samdow @soumith @kshitij12345 @janeyx99

[zou3519]

Here's a minimal repro.

cc @soulitzer @albanD, what do you think the following should return?

import torch
from torch.func import grad

class Cube(torch.autograd.Function):
    @staticmethod
    def forward(x):
        return x ** 3, x

    @staticmethod
    def setup_context(ctx, inputs, outputs):
        cube, x = outputs
        ctx.save_for_backward(x)

    @staticmethod
    def backward(ctx, grad_output, grad_x):
        # NB: grad_x intentionally not used in computation
        x, = ctx.saved_tensors
        result = grad_output * 3 * x ** 2
        return result

def f(x):
    cube, inp = Cube.apply(x)
    return cube

x = torch.tensor(1., requires_grad=True)

print("============= functorch")
grad_x = grad(f)(x)
grad_out = torch.ones_like(grad_x)
result = torch.autograd.grad(grad_x, x, grad_out)
print(result)

print("============= pytorch")
y = f(x)
grad_out = torch.ones_like(y)
grad_x2, = torch.autograd.grad(y, x, create_graph=True)
expected = torch.autograd.grad(grad_x2, x, grad_out)
print(expected)

Spoilers

The question is in the autograd graph (of grad_x and grad_x2), does the x in grad_output * 3 * x ** 2 refer to the output of Cube, or does it refer to the original x?

• If it refers to the original x: then it doesn't matter that we didn't use grad_x in Cube.backward, so the second derivative is computed as 6.
• if it refers to the output of Cube: then it makes sense that the result is 0., because the autograd.Function didn't actually connect the output correctly.

My take is that:

• it should refer to the original x
• the functorch and non-functorch behavior should be consistent (however, I am not sure this is possible to do).

[soulitzer]

Another option is just to error if someone takes an input returns it as an output as-is and then tries to save that for backward through output. The error would tell them to save the input instead. If someone really wanted to return an input as is and then do something fancy with the gradients, e.g. result = grad_output * 3 * x ** 2 + grad_x * 3, then they could just return x.view_as(x) as a workaround.

[albanD]

I would agree that we should not automatically x.alias() when the input is an output if it is being saved for backward.