Repository: pytorch
Issue: Wrapper tensor level confusion
Labels: triaged,needs research,module: __torch_dispatch__,tensor subclass

### 🐛 Describe the bug

With https://github.com/pytorch/pytorch/pull/72832 we have made it possible to go through PyTorch's dispatch stack multiple times. The basic concept is that if you have a situation like WrapperTensor(WrapperTensor(Tensor)), each time you process a wrapper tensor you go through the dispatcher (hitting all of the dispatch keys that are enabled for the wrapper tensor), until you get to the `__torch_dispatch__` call which shunts you back to the top of the dispatch stack for the inner wrapper tensor.

Does this sound similar to what functorch does? It should: functorch without modes operates in much of a similar way, successively processing layered functorch transformations (which you dispatched to from the tensor objects in question) until there's nothing left and you can fall through to more basic PyTorch transformations. However, in functorch, there are two key differences:

* There are only a limited set of dispatch keys which are supported in this way (the ones wedged between FuncTorchDynamicLayerBackMode and FuncTorchDynamicLayerFrontMode)
* Functorch's wrapper tensor also records a *level*, which is used to disambiguate operations between tensors at different levels

Levels are semantically important: suppose that I have BatchedTensor(x, batch_dim=0) + BatchedTensor(y, batch_dim=0); the levels of the two batched tensors tell me disambiguate between these three situations:

* Same level: sum is BatchedTensor(x + y, batch_dim=0)
* x occurs before y (x's BatchedTensor wraps y's): sum is BatchedTensor(BatchedTensor(x.expand(1, b, -1) + y.expand(b, 1, -1), batch_dim=0), batch_dim=0)
* y occurs before x (y's BatchedTensor wraps x's): sum is BatchedTensor(BatchedTensor(x.expand(b, 1, -1) + y.expand(1, b, -1), batch_dim=0), batch_dim=0)

Now, the problem is that, by default, non-functorch wrapper tensors have no level disambiguation (because they were all designed in a universe where there was only ever one level). This is bad, because if a user mixes up two levels you will blindly try to mash the levels together as if they were the same level, and you will get very confusing results.

In previous work (https://fb.quip.com/xvtvA5oJ0b1H), I suggested that we ought to only go through the dispatcher once, and have some concept of "dynamically allocated dispatch keys" (in a lookaside table) for which we could maintain an ordering and use to decide where to go first. Alban's approach is novel, in the sense that he says "whatever, you don't need to keep it all in one go through the dispatch table, we can do as many loops through it as we want to get the levels we want." This suggests a different implementation strategy for handling dispatch:

1. Have every tensor to record an integer level.
2. When doing multiple dispatch, compute the highest level. Discard all tensors which are at lower levels and do multiple dispatch only on tensors at the highest level.

This strategy would be quite inefficient. But here is another way to do it if we just want to optimize the one level case and not the higher level case. Maintain a dispatch key representing tensor "non-zero level", which is at the front of all other keys (this is similar to FuncTorchDynamicLayerFrontMode), and then branch into the alternate dispatch algorithm if you notice you are in this regime. (It is tempting to try a BackendSelect style strategy where this is just a normal dispatch key and the kernel is responsible for calculating the redispatch, but this falls down if the kernel you dispatch to tries to redispatch, thereby bypassing this calculation (it doesn't even work if you don't exclude the higher order key, because redispatch after will skip it). Maybe there is a cute trick that can improve this more.

Assuming we do this, then https://fb.quip.com/xvtvA5oJ0b1H is now just the problem of "how can you define custom C++ classes without adding them to the dispatch table" but we don't have to deal with the level problem. And since we're willing to allocate dispatch keys for all of functorch's basic functionality, it becomes unnecessary for functorch's immediate needs (and we simply need to mainline something like WrapperTensor).

cc @bdhirsh @albanD @zou3519 @Chillee 

### Versions

master