Repository: pytorch
Issue: [export] Running torch._inductor.compile directly on an ExportedProgram fails with non-Tensor inputs
Labels: oncall: pt2,module: inductor,oncall: export

### 🐛 Describe the bug

When trying to export a model that takes non-tensor inputs and running inductor directly on it (to avoid the cost of tracing the model through Dynamo+AOTAutograd again), I get an error about Sympy Boolean types not supported on inductor graph

```python
from pathlib import Path
import torch
import torch._inductor.config
from torch.export import export, Dim, save, load
import torch._dynamo.config

class ToyModel(torch.nn.Module):
    def __init__(self):
        super(ToyModel, self).__init__()
        self.net1 = torch.nn.Embedding(32000, 4096)
        self.net2 = torch.nn.Linear(4096, 32000)
        self.nheads = 32

    def forward(self, x, past_key_value_states=None, use_cache=False):
        y = self.net1(x)
        if use_cache and past_key_value_states is None:
            past_key_value_states = (2 * torch.ones_like(y)).view(*x.shape, self.nheads, -1)
        pkvs_s = past_key_value_states.shape
        y = y * past_key_value_states.view(pkvs_s[0], pkvs_s[1], pkvs_s[2]*pkvs_s[3])
        z = self.net2(y)
        if use_cache:
            return z, past_key_value_states
        return z

toy_model = ToyModel()

torch._inductor.config.fx_graph_cache = True
print("exporting and compiling model")
export_file = Path("model.pt2")
if not export_file.is_file():
    seqlen = Dim("seqlen")
    exported_model = export(
        toy_model,
        (torch.randint(0, 32000, (1, 16)),),
        {"use_cache": True, },
        dynamic_shapes={"x": {1: seqlen}, "use_cache": None})
    save(exported_model, export_file)

# compiling can make first inference pass slow
loaded_model = load(export_file)
compiled_loaded_model = torch._inductor.compile(loaded_model.module(), loaded_model.example_inputs)
```

The error I get is the following:
```
# python scripts/export_repro.py 
exporting and compiling model
Traceback (most recent call last):
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/codecache.py", line 1041, in load
    FxGraphCache._check_can_cache(gm)
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/codecache.py", line 1011, in _check_can_cache
    raise BypassFxGraphCache
torch._inductor.codecache.BypassFxGraphCache

During handling of the above exception, another exception occurred:

Traceback (most recent call last):
  File "/gpfs/users/aviros/foundation-model-stack/scripts/export_repro.py", line 41, in <module>
    compiled_loaded_model = torch._inductor.compile(loaded_model.module(), loaded_model.example_inputs)
                            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/__init__.py", line 29, in compile
    return compile_fx(gm, example_inputs, config_patches=options)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/compile_fx.py", line 1349, in compile_fx
    return handle_dynamo_export_graph(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/compile_fx.py", line 1624, in handle_dynamo_export_graph
    compiled_fn = compile_gm(gm, codegen.process_inputs(*inputs))
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/compile_fx.py", line 1536, in compile_fx
    return aot_autograd(
           ^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/backends/common.py", line 69, in __call__
    cg = aot_module_simplified(gm, example_inputs, **self.kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 974, in aot_module_simplified
    compiled_fn = dispatch_and_compile()
                  ^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 963, in dispatch_and_compile
    compiled_fn, _ = create_aot_dispatcher_function(
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 231, in time_wrapper
    r = func(*args, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 695, in create_aot_dispatcher_function
    compiled_fn, fw_metadata = compiler_fn(
                               ^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 488, in aot_dispatch_autograd
    compiled_fw_func = aot_config.fw_compiler(fw_module, adjusted_flat_args)
                       ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 231, in time_wrapper
    r = func(*args, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/compile_fx.py", line 1441, in fw_compiler_base
    return inner_compile(
           ^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/repro/after_aot.py", line 84, in debug_wrapper
    inner_compiled_fn = compiler_fn(gm, example_inputs)
                        ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/debug.py", line 304, in inner
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 231, in time_wrapper
    r = func(*args, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/compile_fx.py", line 530, in compile_fx_inner
    compiled_graph = FxGraphCache.load(
                     ^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/codecache.py", line 1087, in load
    compiled_graph = compile_fx_fn(gm, example_inputs, **fx_kwargs)
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/compile_fx.py", line 843, in fx_codegen_and_compile
    graph.run(*example_inputs)
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 231, in time_wrapper
    r = func(*args, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/graph.py", line 717, in run
    return super().run(*args)
           ^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/fx/interpreter.py", line 146, in run
    self.env[node] = self.run_node(node)
                     ^^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/graph.py", line 1199, in run_node
    result = super().run_node(n)
             ^^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/fx/interpreter.py", line 203, in run_node
    return getattr(self, n.op)(n.target, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/gpfs/users/aviros/miniforge3/envs/fms-dev/lib/python3.11/site-packages/torch/_inductor/graph.py", line 1062, in output
    assert isinstance(
           ^^^^^^^^^^^
AssertionError: Unsupported inductor graph input type: <class 'sympy.logic.boolalg.BooleanTrue'>
```

### Versions

Collecting environment information...
PyTorch version: 2.5.0.dev20240624+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: CentOS Stream 9 (x86_64)
GCC version: (GCC) 11.4.1 20231218 (Red Hat 11.4.1-3)
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.34

Python version: 3.11.9 | packaged by conda-forge | (main, Apr 19 2024, 18:36:13) [GCC 12.3.0] (64-bit runtime)
Python platform: Linux-4.18.0-372.69.1.el8_6.x86_64-x86_64-with-glibc2.34
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-80GB
GPU 1: NVIDIA A100-SXM4-80GB
GPU 2: NVIDIA A100-SXM4-80GB
GPU 3: NVIDIA A100-SXM4-80GB
GPU 4: NVIDIA A100-SXM4-80GB
GPU 5: NVIDIA A100-SXM4-80GB
GPU 6: NVIDIA A100-SXM4-80GB
GPU 7: NVIDIA A100-SXM4-80GB

Nvidia driver version: 535.104.05
cuDNN version: Could not collect
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

CPU:
Architecture:                    x86_64
CPU op-mode(s):                  32-bit, 64-bit
Address sizes:                   46 bits physical, 57 bits virtual
Byte Order:                      Little Endian
CPU(s):                          80
On-line CPU(s) list:             0-79
Vendor ID:                       GenuineIntel
Model name:                      Intel Xeon Processor (Icelake)
CPU family:                      6
Model:                           134
Thread(s) per core:              2
Core(s) per socket:              20
Socket(s):                       2
Stepping:                        0
BogoMIPS:                        5600.02
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology cpuid tsc_known_freq pni pclmulqdq vmx ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch cpuid_fault invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid avx512f avx512dq rdseed adx smap avx512ifma clflushopt clwb avx512cd sha_ni avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves wbnoinvd arat avx512vbmi umip pku ospke avx512_vbmi2 gfni vaes vpclmulqdq avx512_vnni avx512_bitalg avx512_vpopcntdq la57 rdpid fsrm md_clear arch_capabilities
Virtualization:                  VT-x
Hypervisor vendor:               KVM
Virtualization type:             full
L1d cache:                       2.5 MiB (80 instances)
L1i cache:                       2.5 MiB (80 instances)
L2 cache:                        160 MiB (40 instances)
L3 cache:                        32 MiB (2 instances)
NUMA node(s):                    2
NUMA node0 CPU(s):               0-39
NUMA node1 CPU(s):               40-79
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Retbleed:          Not affected
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Enhanced IBRS, IBPB conditional, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected

Versions of relevant libraries:
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] pytorch-triton==3.0.0+dedb7bdf33
[pip3] torch==2.5.0.dev20240624+cu124
[pip3] torchaudio==2.4.0.dev20240621+cu124
[pip3] torchvision==0.20.0.dev20240621+cu124
[conda] numpy                     1.26.4                   pypi_0    pypi
[conda] pytorch-triton            3.0.0+dedb7bdf33          pypi_0    pypi
[conda] torch                     2.5.0.dev20240624+cu124          pypi_0    pypi
[conda] torchaudio                2.4.0.dev20240621+cu124          pypi_0    pypi
[conda] torchvision               0.20.0.dev20240621+cu124          pypi_0    pypi

cc @ezyang @anijain2305 @chauhang @penguinwu @voznesenskym @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @peterbell10 @ipiszy @yf225 @chenyang78 @kadeng @muchulee8 @ColinPeppler @amjames @desertfire @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4