Repository: pytorch
Issue: Functionalization doesn't work when applied twice with wrapper tensor
Labels: triaged,tensor subclass,module: functionalization

### 🐛 Describe the bug

Apply https://github.com/pytorch/pytorch/pull/73441 and then apply this diff

```
diff --git a/test/test_functionalization.py b/test/test_functionalization.py
index 28476ff2595..d48b5f18d80 100644
--- a/test/test_functionalization.py
+++ b/test/test_functionalization.py
@@ -3,6 +3,7 @@
 import torch
 from torch.testing._internal.common_utils import TestCase, run_tests
 from torch.testing._internal.logging_tensor import LoggingTensor, capture_logs, log_input
+from torch.utils._pytree import tree_map
 
 def are_aliased(x, y):
     if x._base is None and y._base is None:
@@ -14,6 +15,35 @@ def are_aliased(x, y):
     return x._base is y._base
 
 
+class WrapperTensor(torch.Tensor):
+    @staticmethod
+    def __new__(cls, e):
+        r = torch.Tensor._make_wrapper_subclass(cls, e.shape, dtype=e.dtype, requires_grad=False)
+        r.elem = e
+        return r
+
+    __torch_function__ = torch._C._disabled_torch_function_impl
+
+    def __str__(self):
+        return f'WrapperTensor({self.elem})'
+
+    @classmethod
+    def __torch_dispatch__(cls, func, types, args=(), kwargs=None):  # type: ignore
+        def unwrap(e):
+            if isinstance(e, WrapperTensor):
+                return e.elem
+            else:
+                return e
+
+        def wrap(e):
+            if isinstance(e, torch.Tensor):
+                return WrapperTensor(e)
+            else:
+                return e
+
+        return tree_map(wrap, func(*tree_map(unwrap, args), **tree_map(unwrap, kwargs)))
+
+
 class TestFunctionalization(TestCase):
 
     def get_logs(self, func, inpt):
@@ -29,6 +59,13 @@ class TestFunctionalization(TestCase):
                 torch._disable_functionalization()
         return logs
 
+    def test_wrap(self):
+        x = torch.ones(2)
+        x2 = torch._to_functional_tensor(x)
+        x3 = WrapperTensor(x2)
+        x4 = torch._to_functional_tensor(x3)
+        x4 + x4
+
     def assert_functionalization(self, func, inpt):
         input_clone = inpt.clone()
         input_clone2 = inpt.clone()
```

This segfaults. gdb says:

```
#0  0x00007fffddec0cb8 in void c10::function_ref<void (char**, long const*, long, long)>::callback_fn<at::native::AVX2::VectorizedLoop2d<at::native::(anonymous namespace)::add_kernel(at::TensorIteratorBase&, c10::Scalar const&)::{lambda()#2}::operator()() const::{lambda()#8}::operator()() const::{lambda(float, float)#1}, at::native::(anonymous namespace)::add_kernel(at::TensorIteratorBase&, c10::Scalar const&)::{lambda()#2}::operator()() const::{lambda()#2}::operator()() const::{lambda(at::vec::AVX2::Vectorized<float>, at::vec)#2}> >(long, char**, long const*, long, long) ()
   from /data/users/ezyang/pytorch-tmp/torch/lib/libtorch_cpu.so
#1  0x00007fffdb353a3d in at::TensorIteratorBase::serial_for_each(c10::function_ref<void (char**, long const*, long, long)>, at::Range) const ()
   from /data/users/ezyang/pytorch-tmp/torch/lib/libtorch_cpu.so
#2  0x00007fffdb353f67 in at::TensorIteratorBase::for_each(c10::function_ref<void (char**, long const*, long, long)>, long) () from /data/users/ezyang/pytorch-tmp/torch/lib/libtorch_cpu.so
#3  0x00007fffddef89ab in at::native::(anonymous namespace)::add_kernel(at::TensorIteratorBase&, c10::Scalar const&)::{lambda()#2}::operator()() const [clone .isra.8677] ()
   from /data/users/ezyang/pytorch-tmp/torch/lib/libtorch_cpu.so
#4  0x00007fffddef9191 in at::native::(anonymous namespace)::add_kernel(at::TensorIteratorBase&, c10::Scalar const&) () from /data/users/ezyang/pytorch-tmp/torch/lib/libtorch_cpu.so
#5  0x00007fffdb4eb7aa in at::native::structured_add_out::impl(at::Tensor const&, at::Tensor const&, c10::Scalar const&, at::Tensor const&) () from /data/users/ezyang/pytorch-tmp/torch/lib/libtorch_cpu.so
#6  0x00007fffdbe6d83b in at::(anonymous namespace)::wrapper_add_Tensor(at::Tensor const&, at::Tensor const&, c10::Scalar const&) () from /data/users/ezyang/pytorch-tmp/torch/lib/libtorch_cpu.so
#7  0x00007fffdbe6d8c2 in c10::impl::wrap_kernel_functor_unboxed_<c10::impl::detail::WrapFunctionIntoFunctor_<c10::CompileTimeFunctionPointer<at::Tensor (at::Tensor const&, at::Tensor const&, c10::Scalar const&), &at::(anonymous namespace)::wrapper_add_Tensor>, at::Tensor, c10::guts::typelist::typelist<at::Tensor const&, at::Tensor const&, c10::Scalar const&> >, at::Tensor (at::Tensor const&, at::Tensor const&, c10::Scalar const&)>::call(c10::OperatorKernel*, c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, c10::Scalar const&) () from /data/users/ezyang/pytorch-tmp/torch/lib/libtorch_cpu.so
#8  0x00007fffdbc035cc in at::_ops::add_Tensor::redispatch(c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, c10::Scalar const&) () from /data/users/ezyang/pytorch-tmp/torch/lib/libtorch_cpu.so
#9  0x00007fffdcb9f51e in torch::autograd::VariableType::(anonymous namespace)::add_Tensor(c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, c10::Scalar const&) ()
   from /data/users/ezyang/pytorch-tmp/torch/lib/libtorch_cpu.so
#10 0x00007fffdcb9fb5b in c10::impl::make_boxed_from_unboxed_functor<c10::impl::detail::WrapFunctionIntoFunctor_<c10::CompileTimeFunctionPointer<at::Tensor (c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, c10::Scalar const&), &torch::autograd::VariableType::(anonymous namespace)::add_Tensor>, at::Tensor, c10::guts::typelist::typelist<c10::DispatchKeySet, at::Tensor const&, at::Tensor const&, c10::Scalar const&> >, false>::call(c10::OperatorKernel*, c10::OperatorHandle const&, c10::DispatchKeySet, std::vector<c10::IValue, std::allocator<c10::IValue> >*) ()
   from /data/users/ezyang/pytorch-tmp/torch/lib/libtorch_cpu.so
#11 0x00007fffdd6b0550 in c10::Dispatcher::callBoxed(c10::OperatorHandle const&, std::vector<c10::IValue, std::allocator<c10::IValue> >*) const ()
   from /data/users/ezyang/pytorch-tmp/torch/lib/libtorch_cpu.so
#12 0x00007fffe26e511b in torch::jit::invokeOperatorFromPython(std::vector<std::shared_ptr<torch::jit::Operator>, std::allocator<std::shared_ptr<torch::jit::Operator> > > const&, pybind11::args, pybind11::kwargs const&) () from /data/users/ezyang/pytorch-tmp/torch/lib/libtorch_python.so
#13 0x00007fffe26bf345 in torch::jit::initJITBindings(_object*)::{lambda(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&)#136}::operator()(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&) const::{lambda(pybind11::args, {lambda(std::__cxx11::basic_string<char, std::char_traits<char>, std::allocator<char> > const&)#136}::kwargs)#1}::operator()(pybind11, pybind11::args) const ()
   from /data/users/ezyang/pytorch-tmp/torch/lib/libtorch_python.so
```

I'm guessing it has something to do with https://github.com/pytorch/pytorch/issues/73513

### Versions

https://github.com/pytorch/pytorch/pull/73441