Repository: pytorch
Issue: Dynamo export: limited support in Torch.cond
Labels: module: onnx,triaged,oncall: pt2,module: higher order operators

### 🐛 Describe the bug

I'm using Dynamo export on the [DiNTS](https://github.com/Project-MONAI/MONAI/blob/dev/monai/networks/nets/dints.py#L317) model by MONAI. I face an error that conditional statements should be refactored to use Torch.cond(). The documentation of Torch.cond() specifies limitations such as "no inplace mutations on inputs or global variables". Since the model definition relies on input mutations, is there plan to support this in the future? Or would it require a model level refactor? cc @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang @zou3519 @ydwu4 @moraxu

I'm using the following script for export:

```
'''
'''
Install:
pip3 install --pre torch torchvision torchaudio --index-url https://download.pytorch.org/whl/nightly/cu121
pip3 install monai
'''
import torch
from monai.networks.nets import DiNTS, TopologySearch

def loadDiNTS():
    dints_grid_params = {
        "channel_mul": 0.5,
        "num_blocks": 8,
        "num_depths": 4,
        "device": "cuda",
        "use_downsample": False,
        "spatial_dims": 2,
    }
    grid = TopologySearch(**dints_grid_params)

    dints_params = {
        "in_channels": 1,
        "num_classes": 4,
        "act_name": "RELU",
        "norm_name": ("INSTANCE", {"affine": True}),
        "use_downsample": False,
        "spatial_dims": 2,
        "dints_space": grid
    }
    net = DiNTS(**dints_params).to(dints_grid_params["device"])
    return net

model = loadDiNTS()
data = torch.randn(2, 1, 32, 16).to("cuda")
torch.export.export(model, args=(data,))
```
The export error is below

```
Traceback (most recent call last):
  File "/opt/pytorch/repro_123972.py", line 34, in <module>
    torch.export.export(model, args=(data,))
  File "/opt/pytorch/torch/export/__init__.py", line 174, in export
    return _export(
           ^^^^^^^^
  File "/opt/pytorch/torch/export/_trace.py", line 839, in wrapper
    raise e
  File "/opt/pytorch/torch/export/_trace.py", line 822, in wrapper
    ep = fn(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/export/exported_program.py", line 85, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/export/_trace.py", line 1075, in _export
    gm_torch_level = _export_to_torch_ir(
                     ^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/export/_trace.py", line 430, in _export_to_torch_ir
    gm_torch_level, _ = torch._dynamo.export(
                        ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/eval_frame.py", line 1237, in inner
    result_traced = opt_f(*args, **kwargs)
                    ^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/eval_frame.py", line 410, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/nn/modules/module.py", line 1532, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/nn/modules/module.py", line 1541, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/convert_frame.py", line 978, in catch_errors
    return callback(frame, cache_entry, hooks, frame_state, skip=1)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/convert_frame.py", line 411, in _convert_frame_assert
    return _compile(
           ^^^^^^^^^
  File "/opt/pytorch/torch/_utils_internal.py", line 70, in wrapper_function
    return function(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/conda/envs/ptca/lib/python3.11/contextlib.py", line 81, in inner
    return func(*args, **kwds)
           ^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/convert_frame.py", line 700, in _compile
    guarded_code = compile_inner(code, one_graph, hooks, transform)
                   ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/utils.py", line 266, in time_wrapper
    r = func(*args, **kwargs)
        ^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/convert_frame.py", line 568, in compile_inner
    out_code = transform_code_object(code, transform)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/bytecode_transformation.py", line 1116, in transform_code_object
    transformations(instructions, code_options)
  File "/opt/pytorch/torch/_dynamo/convert_frame.py", line 173, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/convert_frame.py", line 515, in transform
    tracer.run()
  File "/opt/pytorch/torch/_dynamo/symbolic_convert.py", line 2237, in run
    super().run()
  File "/opt/pytorch/torch/_dynamo/symbolic_convert.py", line 875, in run
    while self.step():
          ^^^^^^^^^^^
  File "/opt/pytorch/torch/_dynamo/symbolic_convert.py", line 790, in step
    self.dispatch_table[inst.opcode](self, inst)
  File "/opt/pytorch/torch/_dynamo/symbolic_convert.py", line 470, in inner
    raise exc.UserError(
torch._dynamo.exc.UserError: Dynamic control flow is not supported at the moment. Please use functorch.experimental.control_flow.cond to explicitly capture the control flow. For more information about this error, see: https://pytorch.org/docs/main/generated/exportdb/index.html#cond-operands

from user code:
   File "/opt/conda/envs/ptca/lib/python3.11/site-packages/monai/networks/nets/dints.py", line 493, in forward
    if self.node_a[0][d]:

Set TORCH_LOGS="+dynamo" and TORCHDYNAMO_VERBOSE=1 for more information
```

### Versions

Versions of relevant libraries:
[pip3] numpy==1.24.4
[pip3] optree==0.10.0
[pip3] pytorch-quantization==2.1.2
[pip3] pytorch-triton==3.0.0+989adb9a29
[pip3] torch==2.4.0.dev20240412+cu121
[pip3] torch-tensorrt==2.3.0a0
[pip3] torchaudio==2.2.0.dev20240412+cu121
[pip3] torchdata==0.7.1a0
[pip3] torchtext==0.17.0a0
[pip3] torchvision==0.19.0.dev20240412+cu121
[pip3] triton==2.2.0+e28a256

cc @ezyang @msaroufim @bdhirsh @anijain2305 @chauhang