Repository: pytorch
Issue: "TypeError: unhashable type: non-nested SymInt" with `torch.compile`
Labels: triaged,oncall: pt2,module: dynamic shapes,module: aotdispatch,module: pt2-dispatcher

### 🐛 Describe the bug

MWE:

```python
import torch


@torch.compile
def topk(
    A: torch.Tensor,
    B: torch.Tensor,
    A_idx: torch.Tensor,
    B_idx: torch.Tensor,
):
    k_ = min(50, B.size(0))
    dist = torch.cdist(A, B)
    dist.diagonal(int(A_idx - B_idx)).fill_(torch.inf)
    return dist.topk(k_, 1, False)


torch.set_default_device("cuda:0")
A = torch.randn(500, 100)
B = torch.randn(501, 100)
A_idx = torch.tensor(0)
B_idx = torch.tensor(1)
res = topk(A, B, A_idx, B_idx)
B = torch.randn(5, 100)
res = topk(A, B, A_idx, B_idx)
```

`torch.compile` will fail at the last line, i.e., the second `topk` call:

```python
W0904 17:07:39.823000 140480204642112 torch/_dynamo/variables/tensor.py:715] [0/0] Graph break from `Tensor.item()`, consider setting:
W0904 17:07:39.823000 140480204642112 torch/_dynamo/variables/tensor.py:715] [0/0]     torch._dynamo.config.capture_scalar_outputs = True
W0904 17:07:39.823000 140480204642112 torch/_dynamo/variables/tensor.py:715] [0/0] or:
W0904 17:07:39.823000 140480204642112 torch/_dynamo/variables/tensor.py:715] [0/0]     env TORCHDYNAMO_CAPTURE_SCALAR_OUTPUTS=1
W0904 17:07:39.823000 140480204642112 torch/_dynamo/variables/tensor.py:715] [0/0] to include these operations in the captured graph.
W0904 17:07:39.823000 140480204642112 torch/_dynamo/variables/tensor.py:715] [0/0] 
/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_inductor/compile_fx.py:150: UserWarning: TensorFloat32 tensor cores for float32 matrix multiplication available but not enabled. Consider setting `torch.set_float32_matmul_precision('high')` for better performance.
  warnings.warn(
/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/variables/functions.py:663: UserWarning: Graph break due to unsupported builtin None.Tensor.diagonal. This function is either a Python builtin (e.g. _warnings.warn) or a third-party C/C++ Python extension (perhaps created with pybind). If it is a Python builtin, please file an issue on GitHub so the PyTorch team can add support for it and see the next case for a workaround. If it is a third-party C/C++ Python extension, please either wrap it into a PyTorch-understood custom operator (see https://pytorch.org/tutorials/advanced/custom_ops_landing_page.html for more details) or, if it is traceable, use torch.compiler.allow_in_graph.
  torch._dynamo.utils.warn_once(msg)
Traceback (most recent call last):
  File "<private_work_dir>/_topk_compile_test.py", line 25, in <module>
    res = topk(A, B, A_idx, B_idx)
          ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/eval_frame.py", line 433, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "<private_work_dir>/_topk_compile_test.py", line 5, in topk
    @torch.compile
  File "<private_work_dir>/_topk_compile_test.py", line 14, in torch_dynamo_resume_in_topk_at_14
    dist.diagonal(int(A_idx - B_idx)).fill_(torch.inf)
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 1116, in __call__
    return self._torchdynamo_orig_callable(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 948, in __call__
    result = self._inner_convert(
             ^^^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 472, in __call__
    return _compile(
           ^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_utils_internal.py", line 84, in wrapper_function
    return StrobelightCompileTimeProfiler.profile_compile_time(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_strobelight/compile_time_profiler.py", line 129, in profile_compile_time
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 817, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 231, in time_wrapper
    r = func(*args, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 636, in compile_inner
    out_code = transform_code_object(code, transform)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/bytecode_transformation.py", line 1185, in transform_code_object
    transformations(instructions, code_options)
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 178, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/convert_frame.py", line 582, in transform
    tracer.run()
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2451, in run
    super().run()
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 893, in run
    while self.step():
          ^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 805, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2642, in RETURN_VALUE
    self._return(inst)
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/symbolic_convert.py", line 2627, in _return
    self.output.compile_subgraph(
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/output_graph.py", line 1123, in compile_subgraph
    self.compile_and_call_fx_graph(tx, pass2.graph_output_vars(), root)
  File "/root/micromamba/envs/py311/lib/python3.11/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/output_graph.py", line 1318, in compile_and_call_fx_graph
    compiled_fn = self.call_user_compiler(gm)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 231, in time_wrapper
    r = func(*args, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/output_graph.py", line 1409, in call_user_compiler
    raise BackendCompilerFailed(self.compiler_fn, e).with_traceback(
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/output_graph.py", line 1390, in call_user_compiler
    compiled_fn = compiler_fn(gm, self.example_inputs())
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/repro/after_dynamo.py", line 129, in __call__
    compiled_gm = compiler_fn(gm, example_inputs)
                  ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/__init__.py", line 1951, in __call__
    return compile_fx(model_, inputs_, config_patches=self.config)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_inductor/compile_fx.py", line 1505, in compile_fx
    return aot_autograd(
           ^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/backends/common.py", line 69, in __call__
    cg = aot_module_simplified(gm, example_inputs, **self.kwargs)
         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 954, in aot_module_simplified
    compiled_fn, _ = create_aot_dispatcher_function(
                     ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_dynamo/utils.py", line 231, in time_wrapper
    r = func(*args, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 687, in create_aot_dispatcher_function
    compiled_fn, fw_metadata = compiler_fn(
                               ^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 125, in aot_dispatch_base
    flat_fn, flat_args, fw_metadata = pre_compile(
                                      ^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 1939, in pre_compile
    flat_fn, flat_args, fw_metadata = wrapper.pre_compile(
                                      ^^^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 948, in pre_compile
    flat_args_with_synthetic_bases, synthetic_base_info = merge_view_inputs(
                                                          ^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 1336, in merge_view_inputs
    arg_to_old_idx_map = {arg: i for (i, arg) in enumerate(fwd_inputs)}
                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 1336, in <dictcomp>
    arg_to_old_idx_map = {arg: i for (i, arg) in enumerate(fwd_inputs)}
                         ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/root/micromamba/envs/py311/lib/python3.11/site-packages/torch/__init__.py", line 448, in __hash__
    raise TypeError("unhashable type: non-nested SymInt")
torch._dynamo.exc.BackendCompilerFailed: backend='inductor' raised:
TypeError: unhashable type: non-nested SymInt

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
```

If I remove the line `@torch.compile` or `dist.diagonal(...)`, it works well as usual.

### Versions

```python
Collecting environment information...
PyTorch version: 2.4.0
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: TencentOS Server 4.2 (x86_64)
GCC version: (Tencent Compiler 12.3.1.2) 12.3.1 20230912 (TencentOS 12.3.1.2-3)
Clang version: Could not collect
CMake version: version 3.26.5
Libc version: glibc-2.38

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-5.4.241-1-tlinux4-0017.4-x86_64-with-glibc2.38
Is CUDA available: True
CUDA runtime version: 12.2.140
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA A100-SXM4-40GB

Nvidia driver version: 470.239.06
cuDNN version: Probably one of the following:
/usr/local/cuda-12.2/targets/x86_64-linux/lib/libcudnn.so.9
/usr/local/cuda-12.2/targets/x86_64-linux/lib/libcudnn_adv.so.9
/usr/local/cuda-12.2/targets/x86_64-linux/lib/libcudnn_cnn.so.9
/usr/local/cuda-12.2/targets/x86_64-linux/lib/libcudnn_engines_precompiled.so.9
/usr/local/cuda-12.2/targets/x86_64-linux/lib/libcudnn_engines_runtime_compiled.so.9
/usr/local/cuda-12.2/targets/x86_64-linux/lib/libcudnn_graph.so.9
/usr/local/cuda-12.2/targets/x86_64-linux/lib/libcudnn_heuristic.so.9
/usr/local/cuda-12.2/targets/x86_64-linux/lib/libcudnn_ops.so.9
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:                   48 bits physical, 48 bits virtual
Byte Order:                      Little Endian
CPU(s):                          192
On-line CPU(s) list:             0-191
Vendor ID:                       AuthenticAMD
Model name:                      AMD EPYC 7K62 48-Core Processor
CPU family:                      23
Model:                           49
Thread(s) per core:              2
Core(s) per socket:              48
Socket(s):                       2
Stepping:                        0
Frequency boost:                 enabled
CPU(s) scaling MHz:              126%
CPU max MHz:                     2600.0000
CPU min MHz:                     1500.0000
BogoMIPS:                        5190.64
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ht syscall nx mmxext fxsr_opt pdpe1gb rdtscp lm constant_tsc rep_good nopl nonstop_tsc cpuid extd_apicid aperfmperf pni pclmulqdq monitor ssse3 fma cx16 sse4_1 sse4_2 movbe popcnt aes xsave avx f16c rdrand lahf_lm cmp_legacy svm extapic cr8_legacy abm sse4a misalignsse 3dnowprefetch osvw ibs skinit wdt tce topoext perfctr_core perfctr_nb bpext perfctr_llc mwaitx cpb cat_l3 cdp_l3 hw_pstate ssbd mba ibrs ibpb stibp vmmcall fsgsbase bmi1 avx2 smep bmi2 cqm rdt_a rdseed adx smap clflushopt clwb sha_ni xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local clzero irperf xsaveerptr wbnoinvd arat npt lbrv svm_lock nrip_save tsc_scale vmcb_clean flushbyasid decodeassists pausefilter pfthreshold avic v_vmsave_vmload vgif umip rdpid overflow_recov succor smca
Virtualization:                  AMD-V
L1d cache:                       3 MiB (96 instances)
L1i cache:                       3 MiB (96 instances)
L2 cache:                        48 MiB (96 instances)
L3 cache:                        384 MiB (24 instances)
NUMA node(s):                    2
NUMA node0 CPU(s):               0-47,96-143
NUMA node1 CPU(s):               48-95,144-191
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Not affected
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Not affected
Vulnerability Retbleed:          Vulnerable
Vulnerability Spec store bypass: Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:        Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:        Mitigation; Retpolines, IBPB conditional, IBRS_FW, STIBP conditional, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Not affected

Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.4.0
[pip3] torchaudio==2.4.0
[pip3] torchvision==0.19.0
[pip3] triton==3.0.0
[conda] Could not collect
```

cc @ezyang @chauhang @penguinwu @zou3519 @bdhirsh