Repository: pytorch/pytorch

Issue #145761: dynamo cannot trace global op_set .__contains__
Created: 2025-01-27T19:34:18Z
🐛 Describe the bug
This gives a graph break:
import torch

op_set = {
    torch._C._set_grad_enabled,
    torch.amp._enter_autocast,
    torch.amp._exit_autocast,
}

def f(x):
    if torch.ops.aten.add in op_set:
        return x.sin()
    return x.cos()
torch.compile(f, fullgraph=True, backend="eager")(torch.randn(3,4))
error message:
Traceback (most recent call last):
  File "/data/users/yidi/pytorch/test.py", line 13, in <module>
    torch.compile(f, fullgraph=True, backend="eager")(torch.randn(3,4))
  File "/data/users/yidi/pytorch/torch/_dynamo/eval_frame.py", line 566, in _fn
    return fn(*args, **kwargs)
  File "/data/users/yidi/pytorch/torch/_dynamo/convert_frame.py", line 1400, in __call__
    return self._torchdynamo_orig_callable(
  File "/data/users/yidi/pytorch/torch/_dynamo/convert_frame.py", line 565, in __call__
    return _compile(
  File "/data/users/yidi/pytorch/torch/_dynamo/convert_frame.py", line 997, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
  File "/data/users/yidi/pytorch/torch/_utils_internal.py", line 95, in wrapper_function
    return function(*args, **kwargs)
  File "/data/users/yidi/pytorch/torch/_dynamo/convert_frame.py", line 726, in compile_inner
    return _compile_inner(code, one_graph, hooks, transform)
  File "/data/users/yidi/pytorch/torch/_dynamo/convert_frame.py", line 760, in _compile_inner
    out_code = transform_code_object(code, transform)
  File "/data/users/yidi/pytorch/torch/_dynamo/bytecode_transformation.py", line 1404, in transform_code_object
    transformations(instructions, code_options)
  File "/data/users/yidi/pytorch/torch/_dynamo/convert_frame.py", line 236, in _fn
    return fn(*args, **kwargs)
  File "/data/users/yidi/pytorch/torch/_dynamo/convert_frame.py", line 680, in transform
    tracer.run()
  File "/data/users/yidi/pytorch/torch/_dynamo/symbolic_convert.py", line 2914, in run
    super().run()
  File "/data/users/yidi/pytorch/torch/_dynamo/symbolic_convert.py", line 1084, in run
    while self.step():
  File "/data/users/yidi/pytorch/torch/_dynamo/symbolic_convert.py", line 994, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/data/users/yidi/pytorch/torch/_dynamo/symbolic_convert.py", line 2197, in CONTAINS_OP
    self.push(right.call_method(self, "__contains__", [left], {}))
  File "/data/users/yidi/pytorch/torch/_dynamo/variables/user_defined.py", line 818, in call_method
    return super().call_method(tx, name, args, kwargs)
  File "/data/users/yidi/pytorch/torch/_dynamo/variables/base.py", line 413, in call_method
    unimplemented(f"call_method {self} {name} {args} {kwargs}")
  File "/data/users/yidi/pytorch/torch/_dynamo/exc.py", line 361, in unimplemented
    raise Unsupported(msg, case_name=case_name)
torch._dynamo.exc.Unsupported: call_method UserDefinedObjectVariable(set) __contains__ [TorchInGraphFunctionVariable(aten.add)] {}

from user code:
   File "/data/users/yidi/pytorch/test.py", line 10, in f
    if torch.ops.aten.add in op_set:

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information


You can suppress this exception and fall back to eager by setting:
    import torch._dynamo
    torch._dynamo.config.suppress_errors = True

However, when the op_set is a local variable it works fine.
Versions
on master
cc @ezyang @gchanan @zou3519 @kadeng @msaroufim @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @amjames

Duplicate issue #139603: `torch.compile` doesn't support `set` arguments
Created: 2024-11-04T04:01:02Z
🐛 Describe the bug
"""Demonstrate error with `torch.compile` with `set` argument.

pip install torch@https://download.pytorch.org/whl/nightly/cu124/torch-2.6.0.dev20241103%2Bcu124-cp310-cp310-linux_x86_64.whl pytorch-triton@https://download.pytorch.org/whl/nightly/pytorch_triton-3.1.0%2Bcf34004b8a-cp310-cp310-linux_x86_64.whl
"""

import torch

@torch.compile
def func_using_set() -> torch.Tensor:
    with torch.random.fork_rng(devices={torch.device("cuda:0")}):
        return torch.randn((1, 2))

func_using_set()
InternalTorchDynamoError: NotImplementedError: argument of type: <class 'set'>

from user code:
   File "/tmp/ipykernel_765512/1017228397.py", line 6, in func_using_set
    with torch.random.fork_rng(devices={torch.device("cuda:0")}):

cc: @ezyang
Versions
PyTorch version: 2.6.0.dev20241103+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
Versions of relevant libraries:
[pip3] pytorch-triton==3.1.0+cf34004b8a
[pip3] torch==2.6.0.dev20241103+cu124
cc @ezyang @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @amjames

Closing comment (2025-02-27T00:25:24Z): Same as #145761