Repository: pytorch
Issue: Dynamo compiled graph gets overwritten by eager in a data dependent branch when False branch is empty
Labels: triaged,oncall: pt2,module: dynamo,module: guards

### 🐛 Describe the bug

I have encountered an unexpected dynamo capture behavior that is related to a data dependent branch. While the result of the code execution is correct, the way it executes is unexpected -- an already compiled graph got silently dropped and fall back to eager if an empty false branch is run.

It is best illustrated using an example. The example is a simplified version of running HuggingFace Whisper model.

```
import torch
import torch._dynamo as dynamo

global_dict = {
  3: 3
}


def f(x, y):

    # begin graph 0
    x_len = x.shape[-1]
    update_idx = global_dict.get(x_len, None)
    # end graph 0

    if update_idx is not None:
        # begin graph 1
        y[update_idx] += 1
        # end graph 1
    return y


class DebugCompiler:
    def __init__(self):
        self.count = 0

    def __call__(self, gm, example_inputs):
        id = self.count
        print(f'compiling graph {id}')
        gm.graph.print_tabular()

        def run(*args, **kwargs):
            print(f'running graph {id}')
            return gm.forward(*args, **kwargs)

        self.count += 1
        return run


f = dynamo.optimize(DebugCompiler(), dynamic=True)(f)


def test(x_sizes):
    y = torch.zeros(5)
    for size_of_x in x_sizes:
        x = torch.tensor(range(size_of_x))
        y = f(x, y)
        print(y)


# Expected output:
# ```
# compiling graph 0
# running graph 0
# compiling graph 1
# running graph 1
# running graph 0
# running graph 1
# ```
test([3, 3])


# Expected output:
# ```
# running graph 0
# running graph 1
# running graph 0
# running graph 0
# running graph 1
# running graph 0
# running graph 1
# running graph 0
# running graph 1
# ```
# Actual output:
# ```
# running graph 0
# running graph 1
# running graph 0
# running graph 0
# running graph 0
# running graph 0
# ```
test([3, 4, 3, 3, 3])
```

There are two subgraphs in `f`, `graph0` checks the shape of `x` and compare against a dictionary. `graph1` updates `y`. The execution of `graph1` depends on the output of the check of `graph1`.

The unexpected behavior is that once the "False" branch in `f` is triggered, `graph1` is never used. The code for updating `y` silently runs in eager mode. This is proven in the 2nd test case where for the first `3`, graph 1 is used, then the input `4` triggers the False branch, and afterwards, even running with `3` does not trigger the compiled graph anymore. The output `y` for each run is still correct however, which means, the code for updating y is running in eager mode instead of using `graph1`.


On the other hand if there is tensor computation in the False branch this problem does not happen.

```
def f(x, y):
    # begin graph 0
    x_len = x.shape[-1]
    update_idx = global_dict.get(x_len, None)
    # end graph 0
    if update_idx is not None:
        # begin graph 1
        y[update_idx] += 1
        # end graph 1
    else:
       # begin graph 2
        y[update_idx] -= 1
       # end graph 2
    return y
```


### Versions

torch==2.1.0.dev20230331+cu117

cc @ezyang @msaroufim @bdhirsh @anijain2305 @zou3519 @chauhang @voznesenskym @penguinwu @EikanWang @jgong5 @Guobing-Chen @XiaobingSuper @zhuhaozhe @blzheng @wenzhe-nrv @jiayisunx @chenyang78 @kadeng @soumith @wconstab @ngimel