Repository: pytorch
Issue: torch.where(condition) throws an error during export
Labels: triaged,oncall: pt2,module: dynamic shapes,oncall: export

### 🐛 Describe the bug

Using `torch.export` on a module with `torch.where` throws an error. Use the following code snippet to reproduce the error

```python
import torch


class ModuleWithWhere(torch.nn.Module):
    def forward(self, x):
        return torch.where(x > 0)


class ModuleWithNonzero(torch.nn.Module):
    def forward(self, x):
        return torch.nonzero(x > 0, as_tuple=True)


x = torch.rand(2, 10)
mod_nonzero = ModuleWithNonzero()
mod_where = ModuleWithWhere()
dynamic_shapes = ({0: torch.export.Dim("batch_size", max=100)},)

with torch.inference_mode():
    torch.export.export(mod_nonzero, (x,), strict=False, dynamic_shapes=dynamic_shapes)  # passes
    torch.export.export(mod_where, (x,), strict=False, dynamic_shapes=dynamic_shapes)  # fails
```

Error Logs with `TORCH_LOGS="+dynamo"` and `TORCHDYNAMO_VERBOSE=1`
```bash
V1114 06:59:08.253000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:2498] create_env
I1114 06:59:08.427000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:3557] create_symbol s0 = 2 for L['args'][0][0].size()[0] [2, 100] (_export/non_strict_utils.py:109 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s0"
V1114 06:59:08.428000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5358] runtime_assert True == True [statically known]
V1114 06:59:08.429000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5201] eval False == False [statically known]
I1114 06:59:08.431000 2197826 site-packages/torch/_dynamo/utils.py:859] ChromiumEventLogger initialized with id 3870d3ae-2cd4-4175-b365-d460e4974935
V1114 06:59:08.469000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5201] eval False == False [statically known]
V1114 06:59:08.470000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5201] eval True == True [statically known]
V1114 06:59:08.472000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5201] eval Ne(s0, 1) == True [statically known]
V1114 06:59:08.473000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5358] runtime_assert True == True [statically known]
V1114 06:59:08.474000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5358] runtime_assert True == True [statically known]
I1114 06:59:08.476000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:3317] create_unbacked_symint u0 [-int_oo, int_oo] (_subclasses/fake_impls.py:426 in nonzero)
V1114 06:59:08.477000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:4734] _update_var_to_range u0 = VR[0, 9223372036854775806] (update)
I1114 06:59:08.477000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5481] constrain_symbol_range u0 [0, 9223372036854775806]
V1114 06:59:08.479000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5358] runtime_assert u0 >= 0 == True [statically known]
V1114 06:59:08.481000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5358] runtime_assert u0 >= 0 == True [statically known]
V1114 06:59:08.482000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5201] eval Eq(u0, 0) == False [statically known]
V1114 06:59:08.496000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5358] runtime_assert True == True [statically known]
I1114 06:59:08.501000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:3317] create_unbacked_symint u1 [-int_oo, int_oo] (_subclasses/fake_impls.py:426 in nonzero)
V1114 06:59:08.502000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:4734] _update_var_to_range u1 = VR[0, 9223372036854775806] (update)
I1114 06:59:08.502000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5481] constrain_symbol_range u1 [0, 9223372036854775806]
V1114 06:59:08.503000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5358] runtime_assert u1 >= 0 == True [statically known]
V1114 06:59:08.505000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5358] runtime_assert u1 >= 0 == True [statically known]
V1114 06:59:08.506000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5201] eval Eq(u1, 0) == False [statically known]
I1114 06:59:08.509000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:604] compute_unbacked_bindings [u1]
V1114 06:59:08.511000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5201] eval Ne(u1, 1) == True [statically known]
I1114 06:59:08.524000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:3646] produce_guards
V1114 06:59:08.525000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:3830] track_symint L['args'][0][0].size()[0] s0 StrictMinMaxConstraint(warn_only=False, vr=VR[0, 100])
V1114 06:59:08.525000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:3830] track_symint L['args'][0][0].size()[1] 10 None
V1114 06:59:08.525000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:3830] track_symint L['args'][0][0].stride()[0] 10 None
V1114 06:59:08.526000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:3830] track_symint L['args'][0][0].stride()[1] 1 None
V1114 06:59:08.526000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:3830] track_symint L['args'][0][0].storage_offset() 0 None
V1114 06:59:08.535000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5358] runtime_assert u1 >= 0 == True [statically known]
V1114 06:59:08.540000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:2498] create_env
I1114 06:59:08.541000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:3557] create_symbol s0 = 2 for L['args'][0][0].size()[0] [2, 100] (_export/non_strict_utils.py:109 in fakify), for more info run with TORCHDYNAMO_EXTENDED_DEBUG_CREATE_SYMBOL="s0"
V1114 06:59:08.542000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5358] runtime_assert True == True [statically known]
V1114 06:59:08.543000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5201] eval False == False [statically known]
V1114 06:59:08.548000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5201] eval False == False [statically known]
V1114 06:59:08.549000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5201] eval True == True [statically known]
V1114 06:59:08.550000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5201] eval Ne(s0, 1) == True [statically known]
V1114 06:59:08.550000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5358] runtime_assert True == True [statically known]
V1114 06:59:08.551000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5358] runtime_assert True == True [statically known]
I1114 06:59:08.553000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:3317] create_unbacked_symint u0 [-int_oo, int_oo] (_subclasses/fake_impls.py:426 in nonzero)
V1114 06:59:08.554000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:4734] _update_var_to_range u0 = VR[0, 9223372036854775806] (update)
I1114 06:59:08.554000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5481] constrain_symbol_range u0 [0, 9223372036854775806]
V1114 06:59:08.555000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5358] runtime_assert u0 >= 0 == True [statically known]
V1114 06:59:08.556000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5358] runtime_assert u0 >= 0 == True [statically known]
V1114 06:59:08.557000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5201] eval Eq(u0, 0) == False [statically known]
V1114 06:59:08.566000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:5358] runtime_assert True == True [statically known]
Traceback (most recent call last):
  File "/home/lliebenwein/dev/auto-deploy/where_test.py", line 21, in <module>
    torch.export.export(mod_where, (x,), strict=False, dynamic_shapes=dynamic_shapes)  # fails
    ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/export/__init__.py", line 270, in export
    return _export(
           ^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/export/_trace.py", line 1017, in wrapper
    raise e
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/export/_trace.py", line 990, in wrapper
    ep = fn(*args, **kwargs)
         ^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/export/exported_program.py", line 114, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/export/_trace.py", line 1880, in _export
    export_artifact = export_func(  # type: ignore[operator]
                      ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/export/_trace.py", line 1683, in _non_strict_export
    aten_export_artifact = _to_aten_func(  # type: ignore[operator]
                           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/export/_trace.py", line 637, in _export_to_aten_ir
    gm, graph_signature = transform(aot_export_module)(
                          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/export/_trace.py", line 1611, in _aot_export_non_strict
    gm, sig = aot_export(wrapped_mod, args, kwargs=kwargs, **flags)
              ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_functorch/aot_autograd.py", line 1246, in aot_export_module
    fx_g, metadata, in_spec, out_spec = _aot_export_function(
                                        ^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_functorch/aot_autograd.py", line 1480, in _aot_export_function
    fx_g, meta = create_aot_dispatcher_function(
                 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_functorch/aot_autograd.py", line 522, in create_aot_dispatcher_function
    return _create_aot_dispatcher_function(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_functorch/aot_autograd.py", line 759, in _create_aot_dispatcher_function
    compiled_fn, fw_metadata = compiler_fn(
                               ^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/jit_compile_runtime_wrappers.py", line 106, in aot_dispatch_export
    graph, _, _ = aot_dispatch_base_graph(
                  ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py", line 154, in aot_dispatch_base_graph
    fw_module = _create_graph(
                ^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/dispatch_and_compile_graph.py", line 54, in _create_graph
    fx_g = make_fx(
           ^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/fx/experimental/proxy_tensor.py", line 2110, in wrapped
    return make_fx_tracer.trace(f, *args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/fx/experimental/proxy_tensor.py", line 2048, in trace
    return self._trace_inner(f, *args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/fx/experimental/proxy_tensor.py", line 2034, in _trace_inner
    t = dispatch_trace(
        ^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_compile.py", line 32, in inner
    return disable_fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_dynamo/eval_frame.py", line 632, in _fn
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/fx/experimental/proxy_tensor.py", line 1127, in dispatch_trace
    graph = tracer.trace(root, concrete_args)  # type: ignore[arg-type]
            ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/fx/experimental/proxy_tensor.py", line 1631, in trace
    res = super().trace(root, concrete_args)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/fx/_symbolic_trace.py", line 823, in trace
    (self.create_arg(fn(*args)),),
                     ^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/fx/experimental/proxy_tensor.py", line 1182, in wrapped
    out = f(*tensors)
          ^^^^^^^^^^^
  File "<string>", line 1, in <lambda>
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 693, in inner_fn
    outs = fn(*args)
           ^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 413, in _functionalized_f_helper
    f_outs = fn(*f_args)
             ^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 78, in inner_fn
    outs = fn(*args)
           ^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/utils.py", line 182, in flat_fn
    tree_out = fn(*args, **kwargs)
               ^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_functorch/_aot_autograd/traced_function_transforms.py", line 863, in functional_call
    out = mod(*args[params_len:], **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/fx/_symbolic_trace.py", line 801, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/fx/experimental/proxy_tensor.py", line 1701, in call_module
    return Tracer.call_module(self, m, forward, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/fx/_symbolic_trace.py", line 519, in call_module
    ret_val = forward(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/fx/_symbolic_trace.py", line 794, in forward
    return _orig_module_call(mod, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/export/_trace.py", line 1598, in forward
    tree_out = self._export_root(*args, **kwargs)
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/fx/_symbolic_trace.py", line 801, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/fx/experimental/proxy_tensor.py", line 1701, in call_module
    return Tracer.call_module(self, m, forward, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/fx/_symbolic_trace.py", line 519, in call_module
    ret_val = forward(*args, **kwargs)
              ^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/fx/_symbolic_trace.py", line 794, in forward
    return _orig_module_call(mod, *args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1736, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/nn/modules/module.py", line 1747, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/dev/auto-deploy/where_test.py", line 6, in forward
    return torch.where(x > 0)
           ^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/fx/experimental/proxy_tensor.py", line 1230, in __torch_function__
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/fx/experimental/proxy_tensor.py", line 1258, in __torch_function__
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_export/non_strict_utils.py", line 520, in __torch_function__
    return func(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_ops.py", line 833, in handler
    return torch._library.utils.handle_dispatch_mode(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_library/utils.py", line 282, in handle_dispatch_mode
    return curr_mode.__torch_dispatch__(op_overload, overload_types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_subclasses/functional_tensor.py", line 534, in __torch_dispatch__
    outs_unwrapped = func._op_dk(
                     ^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_ops.py", line 833, in handler
    return torch._library.utils.handle_dispatch_mode(
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_library/utils.py", line 282, in handle_dispatch_mode
    return curr_mode.__torch_dispatch__(op_overload, overload_types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/utils/_stats.py", line 21, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/fx/experimental/proxy_tensor.py", line 1308, in __torch_dispatch__
    return proxy_call(self, func, self.pre_dispatch, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/fx/experimental/proxy_tensor.py", line 906, in proxy_call
    out = func(*args, **kwargs)
          ^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_ops.py", line 716, in __call__
    return self._op(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/utils/_stats.py", line 21, in wrapper
    return fn(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_subclasses/fake_tensor.py", line 1238, in __torch_dispatch__
    return self.dispatch(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_subclasses/fake_tensor.py", line 1692, in dispatch
    return self._cached_dispatch_impl(func, types, args, kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_subclasses/fake_tensor.py", line 1348, in _cached_dispatch_impl
    output = self._dispatch_impl(func, types, args, kwargs)
             ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_subclasses/fake_tensor.py", line 1943, in _dispatch_impl
    return decomposition_table[func](*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_prims_common/wrappers.py", line 273, in _fn
    result = fn(*args, **kwargs)
             ^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_prims_common/wrappers.py", line 141, in _fn
    result = fn(**bound.arguments)
             ^^^^^^^^^^^^^^^^^^^^^
  File "/home/lliebenwein/miniconda3/envs/auto/lib/python3.12/site-packages/torch/_refs/__init__.py", line 1926, in where
    raise NotImplementedError
NotImplementedError
I1114 06:59:08.591000 2197826 site-packages/torch/_dynamo/utils.py:399] TorchDynamo compilation metrics:
I1114 06:59:08.591000 2197826 site-packages/torch/_dynamo/utils.py:399] Function                          Runtimes (s)
I1114 06:59:08.591000 2197826 site-packages/torch/_dynamo/utils.py:399] ------------------------------  --------------
I1114 06:59:08.591000 2197826 site-packages/torch/_dynamo/utils.py:399] create_aot_dispatcher_function          0.0919
V1114 06:59:08.591000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:122] lru_cache_stats constrain_symbol_range: CacheInfo(hits=3, misses=6, maxsize=None, currsize=6)
V1114 06:59:08.591000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:122] lru_cache_stats evaluate_expr: CacheInfo(hits=53, misses=12, maxsize=256, currsize=12)
V1114 06:59:08.592000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:122] lru_cache_stats _simplify_floor_div: CacheInfo(hits=0, misses=0, maxsize=None, currsize=0)
V1114 06:59:08.592000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:122] lru_cache_stats _maybe_guard_rel: CacheInfo(hits=0, misses=0, maxsize=256, currsize=0)
V1114 06:59:08.592000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:122] lru_cache_stats _find: CacheInfo(hits=37, misses=5, maxsize=None, currsize=5)
V1114 06:59:08.592000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:122] lru_cache_stats has_hint: CacheInfo(hits=0, misses=0, maxsize=256, currsize=0)
V1114 06:59:08.593000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:122] lru_cache_stats size_hint: CacheInfo(hits=0, misses=0, maxsize=256, currsize=0)
V1114 06:59:08.593000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:122] lru_cache_stats simplify: CacheInfo(hits=6, misses=13, maxsize=None, currsize=13)
V1114 06:59:08.593000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:122] lru_cache_stats _update_divisible: CacheInfo(hits=0, misses=0, maxsize=None, currsize=0)
V1114 06:59:08.593000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:122] lru_cache_stats replace: CacheInfo(hits=1674, misses=60, maxsize=None, currsize=60)
V1114 06:59:08.593000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:122] lru_cache_stats _maybe_evaluate_static: CacheInfo(hits=16, misses=19, maxsize=None, currsize=19)
V1114 06:59:08.594000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:122] lru_cache_stats get_implications: CacheInfo(hits=0, misses=0, maxsize=None, currsize=0)
V1114 06:59:08.594000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:122] lru_cache_stats get_axioms: CacheInfo(hits=10, misses=9, maxsize=None, currsize=9)
V1114 06:59:08.594000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:122] lru_cache_stats safe_expand: CacheInfo(hits=313, misses=35, maxsize=256, currsize=35)
V1114 06:59:08.594000 2197826 site-packages/torch/fx/experimental/symbolic_shapes.py:122] lru_cache_stats uninteresting_files: CacheInfo(hits=38, misses=1, maxsize=None, currsize=1)
```

Note that only the combination of using `torch.inference_mode` and specifying `dynamic_shapes` causes the error. Leaving out either won't throw an error. Writing the same op with `torch.nonzero(..., as_tuple=True)` passes.

I guess it's a known issue though since I can see that there is a `TODO` in the source code to implement `torch.where(condition)` where the error is thrown:  https://github.com/pytorch/pytorch/blob/main/torch/_refs/__init__.py#L1926


### Versions

Collecting environment information...
PyTorch version: 2.5.1+cu124
Is debug build: False
CUDA used to build PyTorch: 12.4
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.6 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.2) 9.4.0
Clang version: Could not collect
CMake version: Could not collect
Libc version: glibc-2.31

Python version: 3.12.5 | packaged by conda-forge | (main, Aug  8 2024, 18:36:51) [GCC 12.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-113-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: Could not collect
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: 
GPU 0: NVIDIA RTX A6000
GPU 1: NVIDIA RTX A6000

Nvidia driver version: 550.54.15
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
Byte Order:                         Little Endian
Address sizes:                      46 bits physical, 48 bits virtual
CPU(s):                             24
On-line CPU(s) list:                0-23
Thread(s) per core:                 2
Core(s) per socket:                 12
Socket(s):                          1
NUMA node(s):                       1
Vendor ID:                          GenuineIntel
CPU family:                         6
Model:                              85
Model name:                         Intel(R) Core(TM) i9-10920X CPU @ 3.50GHz
Stepping:                           7
CPU MHz:                            3402.048
CPU max MHz:                        4800.0000
CPU min MHz:                        1200.0000
BogoMIPS:                           6999.82
Virtualization:                     VT-x
L1d cache:                          384 KiB
L1i cache:                          384 KiB
L2 cache:                           12 MiB
L3 cache:                           19.3 MiB
NUMA node0 CPU(s):                  0-23
Vulnerability Gather data sampling: Mitigation; Microcode
Vulnerability Itlb multihit:        KVM: Mitigation: VMX disabled
Vulnerability L1tf:                 Not affected
Vulnerability Mds:                  Not affected
Vulnerability Meltdown:             Not affected
Vulnerability Mmio stale data:      Mitigation; Clear CPU buffers; SMT vulnerable
Vulnerability Retbleed:             Mitigation; Enhanced IBRS
Vulnerability Spec rstack overflow: Not affected
Vulnerability Spec store bypass:    Mitigation; Speculative Store Bypass disabled via prctl and seccomp
Vulnerability Spectre v1:           Mitigation; usercopy/swapgs barriers and __user pointer sanitization
Vulnerability Spectre v2:           Mitigation; Enhanced IBRS; IBPB conditional; RSB filling; PBRSB-eIBRS SW sequence; BHI Syscall hardening, KVM SW loop
Vulnerability Srbds:                Not affected
Vulnerability Tsx async abort:      Mitigation; TSX disabled
Flags:                              fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush dts acpi mmx fxsr sse sse2 ss ht tm pbe syscall nx pdpe1gb rdtscp lm constant_tsc art arch_perfmon pebs bts rep_good nopl xtopology nonstop_tsc cpuid aperfmperf pni pclmulqdq dtes64 monitor ds_cpl vmx est tm2 ssse3 sdbg fma cx16 xtpr pdcm pcid dca sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand lahf_lm abm 3dnowprefetch cpuid_fault epb cat_l3 cdp_l3 invpcid_single ssbd mba ibrs ibpb stibp ibrs_enhanced tpr_shadow vnmi flexpriority ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 erms invpcid cqm mpx rdt_a avx512f avx512dq rdseed adx smap clflushopt clwb intel_pt avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves cqm_llc cqm_occup_llc cqm_mbm_total cqm_mbm_local dtherm ida arat pln pts hwp hwp_act_window hwp_epp hwp_pkg_req avx512_vnni md_clear flush_l1d arch_capabilities

Versions of relevant libraries:
[pip3] mypy==1.11.2
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.26.4
[pip3] nvidia-cublas-cu12==12.4.5.8
[pip3] nvidia-cuda-cupti-cu12==12.4.127
[pip3] nvidia-cuda-nvrtc-cu12==12.4.127
[pip3] nvidia-cuda-runtime-cu12==12.4.127
[pip3] nvidia-cudnn-cu12==9.1.0.70
[pip3] nvidia-cufft-cu12==11.2.1.3
[pip3] nvidia-curand-cu12==10.3.5.147
[pip3] nvidia-cusolver-cu12==11.6.1.9
[pip3] nvidia-cusparse-cu12==12.3.1.170
[pip3] nvidia-nccl-cu12==2.21.5
[pip3] nvidia-nvjitlink-cu12==12.4.127
[pip3] nvidia-nvtx-cu12==12.4.127
[pip3] torch==2.5.1
[pip3] triton==3.1.0
[conda] numpy                     1.26.4                   pypi_0    pypi
[conda] nvidia-cublas-cu12        12.4.5.8                 pypi_0    pypi
[conda] nvidia-cuda-cupti-cu12    12.4.127                 pypi_0    pypi
[conda] nvidia-cuda-nvrtc-cu12    12.4.127                 pypi_0    pypi
[conda] nvidia-cuda-runtime-cu12  12.4.127                 pypi_0    pypi
[conda] nvidia-cudnn-cu12         9.1.0.70                 pypi_0    pypi
[conda] nvidia-cufft-cu12         11.2.1.3                 pypi_0    pypi
[conda] nvidia-curand-cu12        10.3.5.147               pypi_0    pypi
[conda] nvidia-cusolver-cu12      11.6.1.9                 pypi_0    pypi
[conda] nvidia-cusparse-cu12      12.3.1.170               pypi_0    pypi
[conda] nvidia-nccl-cu12          2.21.5                   pypi_0    pypi
[conda] nvidia-nvjitlink-cu12     12.4.127                 pypi_0    pypi
[conda] nvidia-nvtx-cu12          12.4.127                 pypi_0    pypi
[conda] torch                     2.5.1                    pypi_0    pypi
[conda] triton                    3.1.0                    pypi_0    pypi

cc @ezyang @chauhang @penguinwu @bobrenjc93 @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4