[FX] Cannot deepcopy Tracer objects

[jamesr66a]

🐛 Describe the bug

import torch
import torch.fx

tracer = torch.fx.Tracer()

import copy
tracer_copy = copy.deepcopy(tracer)

---------------------------------------------------------------------------
TypeError                                 Traceback (most recent call last)
[<ipython-input-1-a8c2eaa8e803>](https://localhost:8080/#) in <module>()
      5 
      6 import copy
----> 7 tracer_copy = copy.deepcopy(tracer)

6 frames
[/usr/lib/python3.7/copy.py](https://localhost:8080/#) in deepcopy(x, memo, _nil)
    167                     reductor = getattr(x, "__reduce_ex__", None)
    168                     if reductor:
--> 169                         rv = reductor(4)
    170                     else:
    171                         reductor = getattr(x, "__reduce__", None)

TypeError: can't pickle module objects

But if we override __deepcopy__ we can make this work:

class CopyableTracer(torch.fx.Tracer):
  def __deepcopy__(self, memo):
    new_tracer = CopyableTracer.__new__(CopyableTracer)
    for k, v in self.__dict__.items():
      if k in {'_autowrap_search'}:
        new_obj = copy.copy(v)
      else:
        new_obj = copy.deepcopy(v, memo)

      new_tracer.__dict__[k] = new_obj

    return new_tracer

ct = CopyableTracer()
copied_tracer = copy.deepcopy(ct)

Versions

Collecting environment information... PyTorch version: 1.12.0+cu113 Is debug build: False CUDA used to build PyTorch: 11.3 ROCM used to build PyTorch: N/A

OS: Ubuntu 18.04.6 LTS (x86_64) GCC version: (Ubuntu 7.5.0-3ubuntu1~18.04) 7.5.0 Clang version: 6.0.0-1ubuntu2 (tags/RELEASE_600/final) CMake version: version 3.22.6 Libc version: glibc-2.26

Python version: 3.7.13 (default, Apr 24 2022, 01:04:09) [GCC 7.5.0] (64-bit runtime) Python platform: Linux-5.4.188+-x86_64-with-Ubuntu-18.04-bionic Is CUDA available: False CUDA runtime version: 11.1.105 GPU models and configuration: Could not collect Nvidia driver version: Could not collect cuDNN version: Probably one of the following: /usr/lib/x86_64-linux-gnu/libcudnn.so.7.6.5 /usr/lib/x86_64-linux-gnu/libcudnn.so.8.0.5 /usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.0.5 /usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.0.5 /usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.0.5 /usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.0.5 /usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.0.5 /usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.0.5 HIP runtime version: N/A MIOpen runtime version: N/A Is XNNPACK available: True

Versions of relevant libraries: [pip3] numpy==1.21.6 [pip3] torch==1.12.0+cu113 [pip3] torchaudio==0.12.0+cu113 [pip3] torchsummary==1.5.1 [pip3] torchtext==0.13.0 [pip3] torchvision==0.13.0+cu113 [conda] Could not collect

cc @ezyang @SherlockNoMad