[fake_tensor] torch._subclasses.fake_tensor.DynamicOutputShapeException when calling torch.nonzero using aot_function
Labels: triaged, oncall: pt2, module: dynamic shapes, module: graph breaks

### 🐛 Describe the bug

This issue appears related to https://github.com/pytorch/torchdynamo/issues/1886 but for torch.nonzero instead of torch.repeat_interleave (there are probably others as well).

For some reason using torch._dynamo.optimize works but using aot_function does not. I think this has something to do with needing a graph break but I'm not sure.

It's possible my minimal use case here is too simple; my real use case involves a bunch of computation and then using torch.nonzero to precompute tensor indices that are used later. I'm using aot_function to try and automatically generate code for the forward and backward passes. As a workaround I'm able to split my input into two functions, one that contains everything before the torch.nonzero call and another that takes the resulting indices from the torch.nonzero call as a parameter.

### Error logs

Failed to collect metadata on function, produced code may be suboptimal.  Known situations this can occur are inference mode only compilation involving resize_ or prims (!schema.hasAnyAliasInfo() INTERNAL ASSERT FAILED); if your situation looks different please file a bug to PyTorch.
Traceback (most recent call last):
  File "/torch/_functorch/aot_autograd.py", line 1381, in aot_wrapper_dedupe
    fw_metadata, _out = run_functionalized_fw_and_collect_metadata(flat_fn)(
  File "/torch/_functorch/aot_autograd.py", line 578, in inner
    flat_f_outs = f(*flat_f_args)
  File "/torch/_functorch/aot_autograd.py", line 2314, in flat_fn
    tree_out = fn(*args, **kwargs)
  File "<stdin>", line 2, in f
  File "/torch/utils/_stats.py", line 15, in wrapper
    return fn(*args, **kwargs)
  File "/torch/_subclasses/fake_tensor.py", line 928, in __torch_dispatch__
    op_impl_out = op_impl(self, func, *args, **kwargs)
  File "/torch/_subclasses/fake_tensor.py", line 379, in dyn_shape
    raise DynamicOutputShapeException(func)
torch._subclasses.fake_tensor.DynamicOutputShapeException: aten.nonzero.default

Traceback (most recent call last):
  File "<stdin>", line 1, in <module>
  File "/torch/_functorch/aot_autograd.py", line 2334, in returned_function
    compiled_fn = create_aot_dispatcher_function(
  File "/torch/2.0.0-03c7/lib/python3.10/site-packages/torch/_dynamo/utils.py", line 163, in time_wrapper
    r = func(*args, **kwargs)
  File "/torch/_functorch/aot_autograd.py", line 2184, in create_aot_dispatcher_function
    compiled_fn = compiler_fn(flat_fn, fake_flat_args, aot_config)
  File "/torch/_functorch/aot_autograd.py", line 1504, in aot_wrapper_dedupe
    compiled_fn = compiler_fn(wrapped_flat_fn, deduped_flat_args, aot_config)
  File "/torch/2.0.0-03c7/lib/python3.10/site-packages/torch/_functorch/aot_autograd.py", line 1056, in aot_dispatch_base
    fw_module = make_fx(flat_fn, aot_config.decompositions)(*tmp_flat_args)
  File "/torch/2.0.0-03c7/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 716, in wrapped
    t = dispatch_trace(wrap_key(func, args, fx_tracer), tracer=fx_tracer, concrete_args=tuple(phs))
  File "/torch/2.0.0-03c7/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 209, in _fn
    return fn(*args, **kwargs)
  File "/torch/2.0.0-03c7/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 450, in dispatch_trace
    graph = tracer.trace(root, concrete_args)
  File "/torch/2.0.0-03c7/lib/python3.10/site-packages/torch/_dynamo/eval_frame.py", line 209, in _fn
    return fn(*args, **kwargs)
  File "/torch/fx/_symbolic_trace.py", line 778, in trace
    (self.create_arg(fn(*args)),),
  File "/torch/fx/experimental/proxy_tensor.py", line 466, in wrapped
    out = f(*tensors)
  File "<string>", line 1, in <lambda>
  File "/torch/_functorch/aot_autograd.py", line 1502, in wrapped_flat_fn
    return flat_fn(*add_dupe_args(args))
  File "/torch/_functorch/aot_autograd.py", line 2314, in flat_fn
    tree_out = fn(*args, **kwargs)
  File "<stdin>", line 2, in f
  File "/torch/utils/_stats.py", line 15, in wrapper
    return fn(*args, **kwargs)
  File "/torch/fx/experimental/proxy_tensor.py", line 494, in __torch_dispatch__
    return self.inner_torch_dispatch(func, types, args, kwargs)
  File "/torch/2.0.0-03c7/lib/python3.10/site-packages/torch/fx/experimental/proxy_tensor.py", line 519, in inner_torch_dispatch
    out = proxy_call(self, func, args, kwargs)
  File "/torch/fx/experimental/proxy_tensor.py", line 352, in proxy_call
    out = func(*args, **kwargs)
  File "/torch/_ops.py", line 284, in __call__
    return self._op(*args, **kwargs or {})
  File "/torch/utils/_stats.py", line 15, in wrapper
    return fn(*args, **kwargs)
  File "/torch/_subclasses/fake_tensor.py", line 928, in __torch_dispatch__
    op_impl_out = op_impl(self, func, *args, **kwargs)
  File "/torch/_subclasses/fake_tensor.py", line 379, in dyn_shape
    raise DynamicOutputShapeException(func)
torch._subclasses.fake_tensor.DynamicOutputShapeException: aten.nonzero.default


### Minified repro

import torch
import functorch
from typing import List
import torch._dynamo
from functorch.compile import aot_function, aot_module
def f(x):
  return torch.nonzero(x > 0)

x = torch.ones([4,4])
x[0][3] = 0
x[1][2] = 0
opt_fn = torch._dynamo.optimize("eager")(f)
y = opt_fn(x) #Works

def my_compiler(gm: torch.fx.GraphModule, example_inputs: List[torch.Tensor]):
  print(gm.code)
  return gm.forward

aot_fun = aot_function(f,fw_compiler=my_compiler,bw_compiler=my_compiler)
y1 = aot_fun(x) #Error


### Versions

>>> print(torch.__version__)
2.0.0.dev20230207+cu117
>>> torch.version.git_version
'1530b798ceeff749a7cc5833d9d9627778bd998a'


cc @ezyang @msaroufim @wconstab @bdhirsh @anijain2305 @soumith @ngimel