Repository: pytorch
Issue: __torch_dispatch__ + compile: backward guards
Labels: triaged,module: __torch_dispatch__,oncall: pt2,module: aotdispatch,module: pt2-dispatcher

If you try to run this code:
```
import torch
import torch.utils._pytree as pytree
from torch.utils._python_dispatch import return_and_correct_aliasing

class SubclassTensor(torch.Tensor):
    @staticmethod
    def __new__(cls, a):
        shape = a.shape
        kwargs = {}
        kwargs["strides"] = a.stride()
        kwargs["storage_offset"] = a.storage_offset()
        kwargs["device"] = a.device
        kwargs["layout"] = a.layout
        kwargs["requires_grad"] = a.requires_grad
        kwargs["dtype"] = a.dtype
        out = torch.Tensor._make_wrapper_subclass(cls, shape, **kwargs)
        return out

    __torch_function__ = torch._C._disabled_torch_function_impl

    def __init__(self, a):
        self.a = a

    def __repr__(self):
        a_repr = repr(self.a)
        return f"SubclassTensor({a_repr})"

    def __tensor_flatten__(self):
        return ["a"], None

    @staticmethod
    def __tensor_unflatten__(inner_tensors, meta):
        a = inner_tensors["a"]
        return SubclassTensor(a)

    @classmethod
    def __torch_dispatch__(cls, func, types, args, kwargs):
        if kwargs is None:
            kwargs = {}
        args_a = pytree.tree_map(lambda x: x.a if isinstance(x, SubclassTensor) else x, args)
        kwargs_a = pytree.tree_map(lambda x: x.a if isinstance(x, SubclassTensor) else x, kwargs)
        out_a = func(*args_a, **kwargs_a)
        out = pytree.tree_map(lambda x: SubclassTensor(x) if isinstance(x, torch.Tensor) else x, out_a)
        return return_and_correct_aliasing(func, args, kwargs, out)



@torch.compile
def f(plain, subclass):
    return plain * plain, subclass * subclass


plain = torch.ones(2, requires_grad=True)
subclass = SubclassTensor(torch.ones(2, requires_grad=True))

plain2, subclass2 = f(plain, subclass)
out = plain2 * subclass2
out.sum().backward()
```

You'll run into this error:
```
  File "/data/users/hirsheybar/b/pytorch/torch/_functorch/aot_autograd.py", line 4140, in backward
    assert grad_output_types_ == CompiledFunction.metadata.output_types, f"""\
AssertionError: We incorrectly attempted to compile the backward with incorrect subclass metadata.
If you run into this error, please file an issue.
Expected grad_output types: [<class 'torch.Tensor'>, <class '__main__.SubclassTensor'>]
Got grad_output types: [<class '__main__.SubclassTensor'>, <class '__main__.SubclassTensor'>]
```

**The problem**

What's going on?

(1) When you run torch.compile, AOTAutograd is responsible for tracing through the backward "ahead-of-time", to construct a joint graph (the idea being that we can run a partitioner to decide what to save for the backward vs. recompute, and do all of this while compiling the forward)

(2) In order to trace the backward graph, we need some `grad_output` tensors to pass as input to the backward function that we are tracing

(3) We don't actually know the concrete values of these grad_outputs until the backward executes at runtime. However, we have to make a decision: should the grad_outputs that we create for tracing purposes be subclasses? or plain tensors. This will impact what the backward graph that we trace out looks like.

(4) Today, we make a guess that if your forward outputs are plain tensors then their grad_outs will also be plain tensors. And if your forward outputs are subclasses, then your grad_outs will be the same subclass

That breaks this example Why?

(1) The first output to the compiled function is a plain tensor (`plain * plain`)

(2) However, its `grad_output` will be a `SubclassTensor`. This is because we later ran `(plain2 * subclass2).sum().backward()`, which is what causes subclasses to propagate through the backward graph.

**How do we fix it**

One important point is that even though we traced the backward incorrectly the first time around, we didn't trace it incorrectly in a way that would change any of our partitioning decisions. All nodes in the backward graph that would have changed are "backward-only" nodes, that the partitioner never would have tried to move to the forward graph.

To fix this, we'll need to do something like:

(1) Stash a copy of the (forward) torch IR around in the backward, in case we need to retrace

(2) when the check above fails, instead of erroring, *re-run* our logic that traces out a joint graph and partitions it in a forward backward. This time though, we can pass in (fakeified versions of) the real grad_outs that we witnessed in the backward

(3) use our backward compiler to compile the newly traced backward graph

(4) We now have potentially multiple compiled backward artifacts, for a single compiled forward. We'll need some dispatching logic in the backward that knows which artifact to pick, based on the subclass-ness of the current grad_outs.

Another factor here is compiled autograd. It's worth thinking through how these "backward guards" should work with compiled autograd: if it's easier to support this with compiled autograd, we should consider **only** supporting it with compiled autograd.

cc @Chillee @ezyang @zou3519 @albanD @samdow @msaroufim @wconstab @anijain2305