[PT2] Some errors with `cond` and `torch.compile`
Labels: triaged, oncall: pt2, module: functorch

I am not sure if these are intended to be supported use cases, but as a part of https://github.com/pytorch/pytorch/pull/98775, I experimented with `cond()`. This is not blocking any use case.

```
import torch
from functorch.experimental.control_flow import cond


class Module(torch.nn.Module):
    def __init__(self):
        super().__init__()
        self.linear = torch.nn.Linear(3, 3)

    def forward(self, pred, x):
        def true_fn(val):
            return self.linear(val) * torch.tensor(2)

        def false_fn(val):
            return self.linear(val) * torch.tensor(-1)

        return cond(pred, true_fn, false_fn, [x])

mod = Module()
mod = torch.compile(mod)
x = torch.randn([3, 3])
pred = torch.tensor(x[0][0].item() < 0)
real_result = mod.forward(pred, x)
```
raises
```
File "/fsx/users/andgu/work/pytorch/torch/_ops.py", line 236, in dispatch
    assert final_key in self.py_kernels, f"{dispatch_key} -> {final_key}"
torch._dynamo.exc.BackendCompilerFailed: backend='debug_wrapper' raised:
AssertionError: DispatchKey.Functionalize -> DispatchKey.Functionalize
```

<details>
<summary> Full traceback </summary>

```
Traceback (most recent call last):
  File "dynamo/test_cond.py", line 23, in <module>
    real_result = mod.forward(pred, x)
  File "/fsx/users/andgu/work/pytorch/torch/_dynamo/eval_frame.py", line 118, in forward
    return self.dynamo_ctx(self._orig_mod.forward)(*args, **kwargs)
  File "/fsx/users/andgu/work/pytorch/torch/_dynamo/eval_frame.py", line 247, in _fn
    return fn(*args, **kwargs)
  File "/fsx/users/andgu/work/pytorch/torch/_dynamo/eval_frame.py", line 394, in catch_errors
    return callback(frame, cache_size, hooks)
  File "/fsx/users/andgu/work/pytorch/torch/_dynamo/convert_frame.py", line 453, in _convert_frame
    result = inner_convert(frame, cache_size, hooks)
  File "/fsx/users/andgu/work/pytorch/torch/_dynamo/convert_frame.py", line 113, in _fn
    return fn(*args, **kwargs)
  File "/fsx/users/andgu/work/pytorch/torch/_dynamo/convert_frame.py", line 296, in _convert_frame_assert
    return _compile(
  File "/fsx/users/andgu/work/pytorch/torch/_dynamo/utils.py", line 169, in time_wrapper
    r = func(*args, **kwargs)
  File "/fsx/users/andgu/work/pytorch/torch/_dynamo/convert_frame.py", line 361, in _compile
    out_code = transform_code_object(code, transform)
  File "/fsx/users/andgu/work/pytorch/torch/_dynamo/bytecode_transformation.py", line 683, in transform_code_object
    transformations(instructions, code_options)
  File "/fsx/users/andgu/work/pytorch/torch/_dynamo/convert_frame.py", line 348, in transform
    tracer.run()
  File "/fsx/users/andgu/work/pytorch/torch/_dynamo/symbolic_convert.py", line 1892, in run
    super().run()
  File "/fsx/users/andgu/work/pytorch/torch/_dynamo/symbolic_convert.py", line 611, in run
    and self.step()
  File "/fsx/users/andgu/work/pytorch/torch/_dynamo/symbolic_convert.py", line 571, in step
    getattr(self, inst.opname)(inst)
  File "/fsx/users/andgu/work/pytorch/torch/_dynamo/symbolic_convert.py", line 1979, in RETURN_VALUE
    self.output.compile_subgraph(
  File "/fsx/users/andgu/work/pytorch/torch/_dynamo/output_graph.py", line 630, in compile_subgraph
    self.compile_and_call_fx_graph(tx, list(reversed(stack_values)), root)
  File "/fsx/users/andgu/work/pytorch/torch/_dynamo/output_graph.py", line 700, in compile_and_call_fx_graph
    compiled_fn = self.call_user_compiler(gm)
  File "/fsx/users/andgu/work/pytorch/torch/_dynamo/utils.py", line 169, in time_wrapper
    r = func(*args, **kwargs)
  File "/fsx/users/andgu/work/pytorch/torch/_dynamo/output_graph.py", line 782, in call_user_compiler
    raise BackendCompilerFailed(self.compiler_fn, e).with_traceback(
  File "/fsx/users/andgu/work/pytorch/torch/_dynamo/output_graph.py", line 778, in call_user_compiler
    compiled_fn = compiler_fn(gm, self.fake_example_inputs())
  File "/fsx/users/andgu/work/pytorch/torch/_dynamo/debug_utils.py", line 1098, in debug_wrapper
    compiled_gm = compiler_fn(gm, example_inputs)
  File "/fsx/users/andgu/work/pytorch/torch/__init__.py", line 1530, in __call__
    return compile_fx(model_, inputs_, config_patches=self.config)
  File "/fsx/users/andgu/work/pytorch/torch/_inductor/compile_fx.py", line 722, in compile_fx
    return aot_autograd(
  File "/fsx/users/andgu/work/pytorch/torch/_dynamo/backends/common.py", line 62, in compiler_fn
    cg = aot_module_simplified(gm, example_inputs, **kwargs)
  File "/fsx/users/andgu/work/pytorch/torch/_functorch/aot_autograd.py", line 3093, in aot_module_simplified
    compiled_fn = create_aot_dispatcher_function(
  File "/fsx/users/andgu/work/pytorch/torch/_dynamo/utils.py", line 169, in time_wrapper
    r = func(*args, **kwargs)
  File "/fsx/users/andgu/work/pytorch/torch/_functorch/aot_autograd.py", line 2712, in create_aot_dispatcher_function
    fw_metadata = run_functionalized_fw_and_collect_metadata(
  File "/fsx/users/andgu/work/pytorch/torch/_functorch/aot_autograd.py", line 686, in inner
    flat_f_outs = f(*flat_f_args)
  File "/fsx/users/andgu/work/pytorch/torch/_functorch/aot_autograd.py", line 3017, in functional_call
    out = Interpreter(mod).run(*args[params_len:], **kwargs)
  File "/fsx/users/andgu/work/pytorch/torch/fx/interpreter.py", line 137, in run
    self.env[node] = self.run_node(node)
  File "/fsx/users/andgu/work/pytorch/torch/fx/interpreter.py", line 179, in run_node
    return getattr(self, n.op)(n.target, args, kwargs)
  File "/fsx/users/andgu/work/pytorch/torch/fx/interpreter.py", line 251, in call_function
    return target(*args, **kwargs)
  File "/fsx/users/andgu/work/pytorch/torch/_ops.py", line 247, in __call__
    return torch.overrides.handle_torch_function(
  File "/fsx/users/andgu/work/pytorch/torch/overrides.py", line 1538, in handle_torch_function
    result = mode.__torch_function__(public_api, types, args, kwargs)
  File "/fsx/users/andgu/work/pytorch/torch/_inductor/overrides.py", line 33, in __torch_function__
    return replace_fn(func)(*args, **kwargs)
  File "/fsx/users/andgu/work/pytorch/torch/_ops.py", line 252, in __call__
    return self.dispatch(dispatch_key_set.highestPriorityTypeId(), *args, **kwargs)
  File "/fsx/users/andgu/work/pytorch/torch/_ops.py", line 242, in dispatch
    return kernel(*args, **kwargs)
  File "/fsx/users/andgu/work/pytorch/functorch/experimental/_cond.py", line 122, in cond_autograd
    return cond(pred, true_fn, false_fn, *operands)
  File "/fsx/users/andgu/work/pytorch/torch/_ops.py", line 252, in __call__
    return self.dispatch(dispatch_key_set.highestPriorityTypeId(), *args, **kwargs)
  File "/fsx/users/andgu/work/pytorch/torch/_ops.py", line 236, in dispatch
    assert final_key in self.py_kernels, f"{dispatch_key} -> {final_key}"
torch._dynamo.exc.BackendCompilerFailed: backend='debug_wrapper' raised:
AssertionError: DispatchKey.Functionalize -> DispatchKey.Functionalize

While executing %cond : [#users=1] = call_function[target=torch.ops.cond](args = (%l_pred_, %cond_true_0, %cond_false_0, [%l_x_]), kwargs = {})
Original traceback:
  File "dynamo/test_cond.py", line 17, in forward
    return cond(pred, true_fn, false_fn, [x])
```
</details>

```
import torch
from functorch.experimental.control_flow import cond


x = torch.randn((3,))


def f1(x1, x2):
    return x1 + x2


def f2(x1, x2):
    return x1 * x2


@torch.compile()
def f(z):
    return cond(z, f1, f2, [x, x])

f(torch.tensor(True))
```
raises the same error:
```
File "/fsx/users/andgu/work/pytorch/torch/_ops.py", line 236, in dispatch
    assert final_key in self.py_kernels, f"{dispatch_key} -> {final_key}"
torch._dynamo.exc.BackendCompilerFailed: backend='debug_wrapper' raised:
AssertionError: DispatchKey.Functionalize -> DispatchKey.Functionalize

While executing %cond : [#users=1] = call_function[target=torch.ops.cond](args = (%l_z_, %cond_true_0, %cond_false_0, [%g_x_, %g_x_]), kwargs = {})
Original traceback:
  File "dynamo/test_cond.py", line 18, in f
    return cond(z, f1, f2, [x, x])
```


cc @ezyang @soumith @msaroufim @wconstab @ngimel @bdhirsh @zou3519 @Chillee @samdow @kshitij12345 @janeyx99