autograd.Function should be consistent about returning the same Tensor object if mark_dirty was used.

[zou3519]

🐛 Describe the bug

import torch

class MyRelu(torch.autograd.Function):
    @staticmethod
    def forward(ctx, x):
        result = x.relu_()
        ctx.mark_dirty(result)
        return result

    @staticmethod
    def backward(ctx, grad_output):
        pass

z = torch.tensor(1., requires_grad=True)
x = z.clone()
y = MyRelu.apply(x)
y is x

If z requires_grad, then y is x returns True.
If z does not require grad, then y is x returns False.

This is inconsistent. In-place PyTorch operators always return the same object, no matter the require_grad-ness.

NB: this might be difficult to actually do.

Context

The consistency makes autograd.Function + ctx.mark_dirty with functorch transforms inconsistent.

Versions

main

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