[fake tensor] Deepcopy not supported

[anijain2305]

🐛 Describe the bug

After this #96561, Dynamo exported graphs can't be deepcopied. That PR was added because of an unrelated issue of export not supporting functionalization.

But, I think the fake tensor deepcopy problem still remains. So, opening an issue for it. I am not sure about the priority.

Repro

import torch
import copy
from torch._subclasses import CrossRefFakeMode, FakeTensor, FakeTensorMode

fake_mode = FakeTensorMode()
x = torch.rand(4)
fake_x = fake_mode.from_tensor(x)
print(fake_x)


copy_fake_x = copy.deepcopy(fake_x)
print(copy_fake_x)

Stacktrace

FakeTensor(FakeTensor(..., device='meta', size=(4,)), cpu)
Traceback (most recent call last):
  File "/scratch/anijain/work/pytorch/examples/fake_ops.py", line 11, in <module>
    copy_fake_x = copy.deepcopy(fake_x)
  File "/scratch/anijain/work/env/lib/python3.9/copy.py", line 153, in deepcopy
    y = copier(memo)
  File "/scratch/anijain/work/pytorch/torch/_tensor.py", line 203, in __deepcopy__
    new_tensor.__dict__ = deepcopy(self.__dict__, memo)
  File "/scratch/anijain/work/env/lib/python3.9/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/scratch/anijain/work/env/lib/python3.9/copy.py", line 230, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/scratch/anijain/work/env/lib/python3.9/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/scratch/anijain/work/env/lib/python3.9/copy.py", line 270, in _reconstruct
    state = deepcopy(state, memo)
  File "/scratch/anijain/work/env/lib/python3.9/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/scratch/anijain/work/env/lib/python3.9/copy.py", line 230, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/scratch/anijain/work/env/lib/python3.9/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/scratch/anijain/work/env/lib/python3.9/copy.py", line 270, in _reconstruct
    state = deepcopy(state, memo)
  File "/scratch/anijain/work/env/lib/python3.9/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/scratch/anijain/work/env/lib/python3.9/copy.py", line 230, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/scratch/anijain/work/env/lib/python3.9/copy.py", line 172, in deepcopy
    y = _reconstruct(x, memo, *rv)
  File "/scratch/anijain/work/env/lib/python3.9/copy.py", line 270, in _reconstruct
    state = deepcopy(state, memo)
  File "/scratch/anijain/work/env/lib/python3.9/copy.py", line 146, in deepcopy
    y = copier(x, memo)
  File "/scratch/anijain/work/env/lib/python3.9/copy.py", line 230, in _deepcopy_dict
    y[deepcopy(key, memo)] = deepcopy(value, memo)
  File "/scratch/anijain/work/env/lib/python3.9/copy.py", line 153, in deepcopy
    y = copier(memo)
  File "/scratch/anijain/work/env/lib/python3.9/weakref.py", line 192, in __deepcopy__
    new[deepcopy(key, memo)] = o
  File "/scratch/anijain/work/env/lib/python3.9/copy.py", line 161, in deepcopy
    rv = reductor(4)
TypeError: cannot pickle 'WeakIdRef' object

Versions

NA

[anijain2305]

cc @tugsbayasgalan @eellison

[desertfire]

This does break the AOTInductor flow. There is a simple AOTInductor test I was planning to enable on CI, but it is blocked by this now. @tugsbayasgalan, can we get this fixed soon?

[eellison]

Can we revert the PR? Especially since it is a "hacky" - i.e. not the intended long term solution