pytorch icon indicating copy to clipboard operation
pytorch copied to clipboard

Fix bugs about torch.fx.experimental.proxy_tensor.make_fx

Open FFFrog opened this issue 3 months ago • 4 comments

Stack from ghstack (oldest at bottom):

  • -> #141022

Detailed description:

The codes below will raise an error

import torch
from torch.fx.experimental.proxy_tensor import make_fx

def func(a):
    b = a + 1
    c = b.view(-1)
    c.add_(1)
    return b

input = torch.randn(2)
out = make_fx(func)(input)

The error info are like below:

...
  File "/root/Git.d/pytorch/pytorch/torch/_dynamo/codegen.py", line 34, in <module>
    from .variables.torch_function import TensorWithTFOverrideVariable
  File "/root/Git.d/pytorch/pytorch/torch/_dynamo/variables/torch_function.py", line 185, in <module>
    populate_builtin_to_tensor_fn_map()
  File "/root/Git.d/pytorch/pytorch/torch/_dynamo/variables/torch_function.py", line 146, in populate_builtin_to_tensor_fn_map
    inp0 = torch.ones(1)
  File "/root/Git.d/pytorch/pytorch/torch/fx/experimental/proxy_tensor.py", line 1240, in __torch_function__
    return func(*args, **kwargs)
  File "/root/Git.d/pytorch/pytorch/torch/utils/_stats.py", line 21, in wrapper
    return fn(*args, **kwargs)
  File "/root/Git.d/pytorch/pytorch/torch/fx/experimental/proxy_tensor.py", line 1342, in __torch_dispatch__
    return proxy_call(self, func, self.pre_dispatch, args, kwargs)
  File "/root/Git.d/pytorch/pytorch/torch/fx/experimental/proxy_tensor.py", line 907, in proxy_call
    name=proxy_mode.tracer.graph._target_to_str(func.overloadpacket.__name__),
AttributeError: 'PythonKeyTracer' object has no attribute 'graph'
...

Solutions: Import torch._dynamo before dispatch_trace is called to avoid the context set before dispatch_trace from affecting the torch._dynamo import.

cc @ezyang @SherlockNoMad @EikanWang @jgong5 @wenzhe-nrv

FFFrog avatar Nov 19 '24 14:11 FFFrog