[export/dynamo] torch._check fails at compile time when the condition evaluates to False

[anijain2305]

🐛 Describe the bug

Observed this while debugging an internal model.

import torch

@torch.compile(backend="eager", fullgraph=True)
def fn(x):
    torch._check(x.numel() > 0, lambda: "x must have non-zero elements")
    return x

fn(torch.randn(4))
fn(torch.randn(5))
fn(torch.randn(0))

Error logs

Traceback (most recent call last):
  File "/data/users/anijain/pytorch2/torch/_dynamo/utils.py", line 640, in proxy_args_kwargs
    proxy_args = tuple(arg.as_proxy() for arg in args)
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/users/anijain/pytorch2/torch/_dynamo/utils.py", line 640, in <genexpr>
    proxy_args = tuple(arg.as_proxy() for arg in args)
                       ^^^^^^^^^^^^^^
  File "/data/users/anijain/pytorch2/torch/_dynamo/variables/base.py", line 253, in as_proxy
    raise NotImplementedError(str(self))
NotImplementedError: NestedUserFunctionVariable()

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/data/users/anijain/pytorch2/examples/ketan.py", line 8, in <module>
    fn(torch.randn(4))
  File "/data/users/anijain/pytorch2/torch/_dynamo/eval_frame.py", line 434, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/data/users/anijain/pytorch2/torch/_dynamo/convert_frame.py", line 1121, in __call__
    return self._torchdynamo_orig_callable(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/users/anijain/pytorch2/torch/_dynamo/convert_frame.py", line 472, in __call__
    return _compile(
           ^^^^^^^^^
  File "/data/users/anijain/pytorch2/torch/_utils_internal.py", line 84, in wrapper_function
    return StrobelightCompileTimeProfiler.profile_compile_time(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/users/anijain/pytorch2/torch/_strobelight/compile_time_profiler.py", line 129, in profile_compile_time
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/anijain/local/miniconda3/envs/pytorch2/lib/python3.11/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/data/users/anijain/pytorch2/torch/_dynamo/convert_frame.py", line 817, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/users/anijain/pytorch2/torch/_dynamo/utils.py", line 233, in time_wrapper
    r = func(*args, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^
  File "/data/users/anijain/pytorch2/torch/_dynamo/convert_frame.py", line 636, in compile_inner
    out_code = transform_code_object(code, transform)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/users/anijain/pytorch2/torch/_dynamo/bytecode_transformation.py", line 1270, in transform_code_object
    transformations(instructions, code_options)
  File "/data/users/anijain/pytorch2/torch/_dynamo/convert_frame.py", line 178, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/data/users/anijain/pytorch2/torch/_dynamo/convert_frame.py", line 582, in transform
    tracer.run()
  File "/data/users/anijain/pytorch2/torch/_dynamo/symbolic_convert.py", line 2476, in run
    super().run()
  File "/data/users/anijain/pytorch2/torch/_dynamo/symbolic_convert.py", line 904, in run
    while self.step():
          ^^^^^^^^^^^
  File "/data/users/anijain/pytorch2/torch/_dynamo/symbolic_convert.py", line 816, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/data/users/anijain/pytorch2/torch/_dynamo/symbolic_convert.py", line 510, in wrapper
    return inner_fn(self, inst)
           ^^^^^^^^^^^^^^^^^^^^
  File "/data/users/anijain/pytorch2/torch/_dynamo/symbolic_convert.py", line 2083, in CALL
    self.call_function(fn, args, kwargs)
  File "/data/users/anijain/pytorch2/torch/_dynamo/symbolic_convert.py", line 754, in call_function
    self.push(fn.call_function(self, args, kwargs))
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/users/anijain/pytorch2/torch/_dynamo/variables/torch.py", line 781, in call_function
    *proxy_args_kwargs(args, kwargs),
     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/data/users/anijain/pytorch2/torch/_dynamo/utils.py", line 647, in proxy_args_kwargs
    unimplemented(
  File "/data/users/anijain/pytorch2/torch/_dynamo/exc.py", line 220, in unimplemented
    raise Unsupported(msg) from from_exc
torch._dynamo.exc.Unsupported: call_function args: ConstantVariable(bool: True) NestedUserFunctionVariable()

from user code:
   File "/data/users/anijain/pytorch2/examples/ketan.py", line 5, in fn
    torch._check(x.numel() > 0, lambda: "x must have non-zero elements")

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

Minified repro

No response

Versions

N/A

cc @ezyang @chauhang @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4

[NiharJani2002]

Hiii, I am not official contributor but I have gone through this issue, and also experimented with the code. It simply says dynamo is not supported for python 3.12 version and above.

[NiharJani2002]

[NiharJani2002]

Either you should downgrade your python version from 3.12 or the community should start making changes in the repo upgrading the python version to 3.12 and above for dynamo support.

[NiharJani2002]

If I am community is deciding to make changes for dynamo support for python version 3.12 and above I am eager to contribute in this !!!!.

[masnesral]

Just based on the title, not sure if this belongs to dynamo or export folks. I'll kick it to export for now and assume the export oncall knows more than I do about how to categorize this one :/

[anijain2305]

@NiharJani2002 I think you have a old version of PyTorch. nightly supports 3.12.

[NiharJani2002]

@anijain2305. No I am using latest stable version of Pytorch which is 2.3.1