Repository: pytorch
Issue: torch.cond with shape condition in predicate doesn't allow for assertions
Labels: triaged,oncall: pt2,module: dynamic shapes,module: higher order operators,module: pt2-dispatcher

### 🐛 Describe the bug

```python
import torch
from torch.nn.attention.flex_attention import create_block_mask, flex_attention
torch.set_default_device("cuda")

@torch.compile(dynamic=True)
def f(x):
    def g(x):
        return x.view(x.shape[0] * x.shape[1])

    def h(x):
        assert x.shape[0] == 1
        return x.view(x.shape[1]) * 2

    return torch.cond(x.shape[0] > 1, g, h, (x,))

f(torch.randn(8, 12))
```

In this case, I would expect this to be able to run correctly. But it errors with
```
Traceback (most recent call last):
  File "/home/chilli/local/pytorch/torch/_dynamo/variables/higher_order_ops.py", line 51, in graph_break_as_hard_error
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/variables/higher_order_ops.py", line 756, in call_function
    ) = speculate_branch(False)
        ^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/variables/higher_order_ops.py", line 730, in speculate_branch
    ) = speculate_subgraph(
        ^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/variables/higher_order_ops.py", line 548, in speculate_subgraph
    raise ex
  File "/home/chilli/local/pytorch/torch/_dynamo/variables/higher_order_ops.py", line 477, in speculate_subgraph
    output = f.call_function(tx, args, sub_kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/variables/functions.py", line 111, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/symbolic_convert.py", line 948, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/symbolic_convert.py", line 3103, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/symbolic_convert.py", line 3231, in inline_call_
    tracer.run()
  File "/home/chilli/local/pytorch/torch/_dynamo/symbolic_convert.py", line 1095, in run
    while self.step():
          ^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/symbolic_convert.py", line 1007, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/home/chilli/local/pytorch/torch/_dynamo/symbolic_convert.py", line 485, in inner
    unimplemented(
  File "/home/chilli/local/pytorch/torch/_dynamo/exc.py", line 304, in unimplemented
    raise Unsupported(msg, case_name=case_name)
torch._dynamo.exc.Unsupported: Assertion failed on symbolic shapes. Did you make sure eager mode succeeds?

The above exception was the direct cause of the following exception:

Traceback (most recent call last):
  File "/data/users/chilli/t.py", line 16, in <module>
    f(torch.randn(8, 12))
  File "/home/chilli/local/pytorch/torch/_dynamo/eval_frame.py", line 487, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/convert_frame.py", line 1364, in __call__
    return self._torchdynamo_orig_callable(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/convert_frame.py", line 1155, in __call__
    result = self._inner_convert(
             ^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/convert_frame.py", line 544, in __call__
    return _compile(
           ^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/convert_frame.py", line 964, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/convert_frame.py", line 695, in compile_inner
    return _compile_inner(code, one_graph, hooks, transform)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_utils_internal.py", line 87, in wrapper_function
    return function(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/convert_frame.py", line 728, in _compile_inner
    out_code = transform_code_object(code, transform)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/bytecode_transformation.py", line 1337, in transform_code_object
    transformations(instructions, code_options)
  File "/home/chilli/local/pytorch/torch/_dynamo/convert_frame.py", line 229, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/convert_frame.py", line 657, in transform
    tracer.run()
  File "/home/chilli/local/pytorch/torch/_dynamo/symbolic_convert.py", line 2888, in run
    super().run()
  File "/home/chilli/local/pytorch/torch/_dynamo/symbolic_convert.py", line 1095, in run
    while self.step():
          ^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/symbolic_convert.py", line 1007, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/home/chilli/local/pytorch/torch/_dynamo/symbolic_convert.py", line 615, in wrapper
    return inner_fn(self, inst)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/symbolic_convert.py", line 2395, in CALL
    self._call(inst)
  File "/home/chilli/local/pytorch/torch/_dynamo/symbolic_convert.py", line 2389, in _call
    self.call_function(fn, args, kwargs)
  File "/home/chilli/local/pytorch/torch/_dynamo/symbolic_convert.py", line 942, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/variables/functions.py", line 339, in call_function
    return super().call_function(tx, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/variables/functions.py", line 111, in call_function
    return tx.inline_user_function_return(self, [*self.self_args(), *args], kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/symbolic_convert.py", line 948, in inline_user_function_return
    return InliningInstructionTranslator.inline_call(self, fn, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/symbolic_convert.py", line 3103, in inline_call
    return cls.inline_call_(parent, func, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/symbolic_convert.py", line 3231, in inline_call_
    tracer.run()
  File "/home/chilli/local/pytorch/torch/_dynamo/symbolic_convert.py", line 1095, in run
    while self.step():
          ^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/symbolic_convert.py", line 1007, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/home/chilli/local/pytorch/torch/_dynamo/symbolic_convert.py", line 615, in wrapper
    return inner_fn(self, inst)
           ^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/symbolic_convert.py", line 2395, in CALL
    self._call(inst)
  File "/home/chilli/local/pytorch/torch/_dynamo/symbolic_convert.py", line 2389, in _call
    self.call_function(fn, args, kwargs)
  File "/home/chilli/local/pytorch/torch/_dynamo/symbolic_convert.py", line 942, in call_function
    self.push(fn.call_function(self, args, kwargs))  # type: ignore[arg-type]
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/chilli/local/pytorch/torch/_dynamo/variables/higher_order_ops.py", line 54, in graph_break_as_hard_error
    raise UncapturedHigherOrderOpError(reason + msg) from e
torch._dynamo.exc.UncapturedHigherOrderOpError: Cond doesn't work unless it is captured completely with torch.compile. Scroll up to find out what causes the graph break.

from user code:
   File "/data/users/chilli/t.py", line 14, in f
    return torch.cond(x.shape[0] > 1, g, h, (x,))
  File "/home/chilli/local/pytorch/torch/_higher_order_ops/cond.py", line 131, in cond
    return cond_op(pred, true_fn, false_fn, operands)

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
```

### Versions

N/A

cc @ezyang @chauhang @penguinwu @bobrenjc93 @zou3519 @ydwu4 @bdhirsh @yf225