Error when torch.export.export a dynamic shape module

HolyWu: <p>Example to reproduce:</p>
<pre><code class="lang-auto">import torch


class MyModule(torch.nn.Module):
    def __init__(self):
        super().__init__()

    def forward(self, x, y):
        x = torch.nn.functional.interpolate(x, scale_factor=0.5, mode="bilinear")
        x = torch.nn.functional.interpolate(x, scale_factor=2.0, mode="bilinear")
        x = x + y
        return x


device = torch.device("cuda", 0)
model = MyModule().eval().to(device)

inputs = (torch.rand((1, 1, 48, 48), device=device), torch.rand((1, 1, 48, 48), device=device))

dim = torch.export.Dim("Dim", min=16, max=128)
dynamic_shapes = {"x": {2: dim, 3: dim}, "y": {2: dim, 3: dim}}

exported_program = torch.export.export(model, inputs, dynamic_shapes=dynamic_shapes)
print(exported_program)
</code></pre>
<p>Got the following error:</p>
<pre><code class="lang-auto">Traceback (most recent call last):
  File "/home/holywu/test.py", line 23, in &lt;module&gt;
    exported_program = torch.export.export(model, inputs, dynamic_shapes=dynamic_shapes)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/export/__init__.py", line 174, in export
    return _export(
           ^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/export/_trace.py", line 1071, in wrapper
    raise e
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/export/_trace.py", line 1043, in wrapper
    ep = fn(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/export/exported_program.py", line 100, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/export/_trace.py", line 2039, in _export
    export_artifact = export_func(  # type: ignore[operator]
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/export/_trace.py", line 1278, in _strict_export
    return _strict_export_lower_to_aten_ir(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/export/_trace.py", line 1306, in _strict_export_lower_to_aten_ir
    gm_torch_level = _export_to_torch_ir(
                     ^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/export/_trace.py", line 552, in _export_to_torch_ir
    gm_torch_level, _ = torch._dynamo.export(
                        ^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 1410, in inner
    result_traced = opt_f(*args, **kwargs)
                    ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1735, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1746, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 462, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1735, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1746, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 1182, in __call__
    return self._torchdynamo_orig_callable(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 506, in __call__
    return _compile(
           ^^^^^^^^^
  File "/usr/lib/python3.12/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 889, in _compile
    raise InternalTorchDynamoError(str(e)).with_traceback(
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 862, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/_dynamo/utils.py", line 272, in time_wrapper
    r = func(*args, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/_utils_internal.py", line 85, in wrapper_function
    return StrobelightCompileTimeProfiler.profile_compile_time(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/_strobelight/compile_time_profiler.py", line 129, in profile_compile_time
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 680, in compile_inner
    out_code = transform_code_object(code, transform)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/_dynamo/bytecode_transformation.py", line 1322, in transform_code_object
    transformations(instructions, code_options)
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 201, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/_dynamo/convert_frame.py", line 622, in transform
    tracer.run()
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 2657, in run
    super().run()
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 925, in run
    while self.step():
          ^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 837, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 2848, in RETURN_VALUE
    self._return(inst)
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/_dynamo/symbolic_convert.py", line 2833, in _return
    self.output.compile_subgraph(
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/_dynamo/output_graph.py", line 1109, in compile_subgraph
    self.compile_and_call_fx_graph(tx, list(reversed(stack_values)), root)
  File "/usr/lib/python3.12/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/_dynamo/output_graph.py", line 1299, in compile_and_call_fx_graph
    insert_deferred_runtime_asserts(
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/fx/passes/runtime_assert.py", line 288, in insert_deferred_runtime_asserts
    add_runtime_asserts(ras_by_symbol.pop(None, []))  # type: ignore[call-overload]
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/fx/passes/runtime_assert.py", line 221, in add_runtime_asserts
    res = _sympy_interp(expr_to_proxy, ra.expr).node
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/fx/passes/runtime_assert.py", line 179, in _sympy_interp
    [_sympy_interp(expr_to_proxy, arg) for arg in expr.args],
     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/fx/passes/runtime_assert.py", line 179, in _sympy_interp
    [_sympy_interp(expr_to_proxy, arg) for arg in expr.args],
     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/fx/passes/runtime_assert.py", line 177, in _sympy_interp
    expr_to_proxy[expr] = _run_sympy_handler(
                          ^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/utils/_sympy/interp.py", line 134, in _run_sympy_handler
    return getattr(analysis, handler_name)(*args, index_dtype)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/utils/_sympy/reference.py", line 87, in trunc_to_int
    return TruncToInt(x)
           ^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/sympy/core/cache.py", line 72, in wrapper
    retval = cfunc(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/sympy/core/function.py", line 466, in __new__
    result = super().__new__(cls, *args, **options)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/sympy/core/cache.py", line 72, in wrapper
    retval = cfunc(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/sympy/core/function.py", line 307, in __new__
    evaluated = cls.eval(*args)
                ^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/utils/_sympy/functions.py", line 652, in eval
    if number in (sympy.oo, int_oo):
       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/fx/proxy.py", line 450, in __bool__
    return self.tracer.to_bool(self)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/holywu/.local/lib/python3.12/site-packages/torch/fx/proxy.py", line 310, in to_bool
    raise TraceError('symbolically traced variables cannot be used as inputs to control flow')
torch._dynamo.exc.InternalTorchDynamoError: symbolically traced variables cannot be used as inputs to control flow

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information
</code></pre>
<p>Logs with TORCH_LOGS=“+dynamo” and TORCHDYNAMO_VERBOSE=1: <a href="https://gist.github.com/HolyWu/62070659d4a22a7c4c45e8b6202b274f" class="inline-onebox" rel="noopener nofollow ugc">gist:62070659d4a22a7c4c45e8b6202b274f · GitHub</a></p>