Repository: pytorch
Issue: Comparing dynamic shapes fails with KeyError
Labels: triaged,oncall: pt2,module: dynamic shapes

### 🐛 Describe the bug

Unfortunately, I am not able to provide a minimum example.

In our custom backend, we check if the shape that goes through a module changes. 
So we are storing the shape on the module, and if there's a re-compilation happen, we check if that shape changed by comparing the old stored shape with the new shape.

I could imagine this bug happens if the old stored shape has references to symbols that the recompiled graph does not have.

It's not a of high priority as we are doing admittedly something unusual, but wanted to report.

### Error logs

```
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/transformers/generation/utils.py:1527: in <resume in generate>
    result = self._greedy_search(
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/transformers/generation/utils.py:2370: in _greedy_search
    eos_token_id_tensor = torch.tensor(eos_token_id).to(input_ids.device) if eos_token_id is not None else None
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/transformers/generation/utils.py:2411: in <resume in _greedy_search>
    outputs = self(
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/nn/modules/module.py:1518: in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/nn/modules/module.py:1527: in _call_impl
    return forward_call(*args, **kwargs)
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/_dynamo/eval_frame.py:490: in catch_errors
    return callback(frame, cache_entry, hooks, frame_state)
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py:641: in _convert_frame
    result = inner_convert(frame, cache_size, hooks, frame_state)
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py:133: in _fn
    return fn(*args, **kwargs)
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py:389: in _convert_frame_assert
    return _compile(
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py:569: in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/_dynamo/utils.py:189: in time_wrapper
    r = func(*args, **kwargs)
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py:491: in compile_inner
    out_code = transform_code_object(code, transform)
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/_dynamo/bytecode_transformation.py:1028: in transform_code_object
    transformations(instructions, code_options)
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/_dynamo/convert_frame.py:458: in transform
    tracer.run()
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py:2069: in run
    super().run()
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py:719: in run
    and self.step()
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py:683: in step
    getattr(self, inst.opname)(inst)
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/_dynamo/symbolic_convert.py:2157: in RETURN_VALUE
    self.output.compile_subgraph(
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/_dynamo/output_graph.py:857: in compile_subgraph
    self.compile_and_call_fx_graph(tx, pass2.graph_output_vars(), root)
/opt/pyenv-root/versions/3.9.17/lib/python3.9/contextlib.py:79: in inner
    return func(*args, **kwds)
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/_dynamo/output_graph.py:957: in compile_and_call_fx_graph
    compiled_fn = self.call_user_compiler(gm)
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/_dynamo/utils.py:189: in time_wrapper
    r = func(*args, **kwargs)
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/_dynamo/output_graph.py:1024: in call_user_compiler
    raise BackendCompilerFailed(self.compiler_fn, e).with_traceback(
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/_dynamo/output_graph.py:1009: in call_user_compiler
    compiled_fn = compiler_fn(gm, self.example_inputs())
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/_dynamo/repro/after_dynamo.py:117: in debug_wrapper
    compiled_gm = compiler_fn(gm, example_inputs)
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/recogni/torch/conversion/prepare/_dynamo.py:297: in _backend
    if old_input_shape != input_shape:
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/__init__.py:352: in __bool__
    return self.node.bool_()
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/fx/experimental/symbolic_shapes.py:972: in bool_
    return self.guard_bool("", 0)
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/fx/experimental/symbolic_shapes.py:954: in guard_bool
    r = self.shape_env.evaluate_expr(self.expr, self.hint, fx_node=self.fx_node)
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/fx/experimental/symbolic_shapes.py:3500: in evaluate_expr
    static_expr = self._maybe_evaluate_static(expr)
/opt/pyenv-root/versions/3.9.17/lib/python3.9/site-packages/torch/fx/experimental/symbolic_shapes.py:1505: in wrapper
    return fn_cache(self, *args, **kwargs)

self = <torch.fx.experimental.symbolic_shapes.ShapeEnv object at 0x7ff44de9c130>
expr = Eq(s0, s1 + 1)

    @_lru_cache
    def _maybe_evaluate_static(
        self, expr: "sympy.Expr", *, unbacked_only: bool = False, compute_hint: bool = False
    ) -> "Optional[sympy.Expr]":
        """
        Tries to evaluate expr without introducing guards
        If unbacked_only == True, then we only do substitutions on
        unbacked SymInts (leaving regular hinted integers alone).
        """
        expr = self.simplify(expr)
        symbols = list(expr.free_symbols)
        # Apply known runtime asserts
        for s in symbols:
            # Unbacked symints only
            if s in self.var_to_val:
                continue
            subst = {}
            for ra in self.deferred_runtime_asserts[s]:
                if compute_hint:
                    e = ra.expr.xreplace(self.var_to_val)
                else:
                    e = ra.expr
                subst[e] = sympy.true
                subst[sympy.Not(e)] = sympy.false
                # NB: this doesn't match relations if they're flipped; e.g.,
                # if you have x < 5, we won't get 5 > x.  Holler if this is
                # a problem
            # NB: this helps us deal with And/Or connectives
            expr = expr.subs(subst)
        # Simplify making use of value range lower bound
        new_shape_env = {}
        new_range_env = {}
        for idx, k in enumerate(symbols):
>           vr = self.var_to_range[k]
E           torch._dynamo.exc.BackendCompilerFailed: backend='_backend' raised:
E           KeyError: s1
E           
E           Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information
E           
E           
E           You can suppress this exception and fall back to eager by setting:
E               import torch._dynamo
E               torch._dynamo.config.suppress_errors = True
```

### Minified repro

_No response_

### Versions

```
PyTorch version: 2.1.1+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.2 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: Could not collect
CMake version: version 3.28.3
Libc version: glibc-2.35

Python version: 3.10.6 (main, May 29 2023, 11:10:38) [GCC 11.3.0] (64-bit runtime)
Python platform: Linux-5.19.0-1030-gcp-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 11.8.89
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: Tesla T4
Nvidia driver version: 550.54.14
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, 48 bits virtual
Byte Order:                      Little Endian
CPU(s):                          8
On-line CPU(s) list:             0-7
Vendor ID:                       GenuineIntel
Model name:                      Intel(R) Xeon(R) CPU @ 2.20GHz
CPU family:                      6
Model:                           79
Thread(s) per core:              2
Core(s) per socket:              4
Socket(s):                       1
Stepping:                        0
BogoMIPS:                        4399.99
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 nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single pti ssbd ibrs ibpb stibp fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm rdseed adx smap xsaveopt arat md_clear arch_capabilities
Hypervisor vendor:               KVM
Virtualization type:             full
L1d cache:                       128 KiB (4 instances)
L1i cache:                       128 KiB (4 instances)
L2 cache:                        1 MiB (4 instances)
L3 cache:                        55 MiB (1 instance)
NUMA node(s):                    1
NUMA node0 CPU(s):               0-7
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Mitigation; PTE Inversion
Vulnerability Mds:               Mitigation; Clear CPU buffers; SMT Host state unknown
Vulnerability Meltdown:          Mitigation; PTI
Vulnerability Mmio stale data:   Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Retbleed:          Mitigation; IBRS
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; IBRS, IBPB conditional, STIBP conditional, RSB filling, PBRSB-eIBRS Not affected
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Mitigation; Clear CPU buffers; SMT Host state unknown
```

cc @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang