Repository: pytorch
Issue: Cannot use constrain_as_size from fake tensor implementations: RuntimeError: tried to get Int out of SymInt
Labels: good first issue,triaged,oncall: pt2,module: dynamic shapes,oncall: export

### 🐛 Describe the bug

Patch in:

```
diff --git a/torch/_subclasses/fake_tensor.py b/torch/_subclasses/fake_tensor.py
index 34408a9b3f4..00cfce23df0 100644
--- a/torch/_subclasses/fake_tensor.py
+++ b/torch/_subclasses/fake_tensor.py
@@ -585,7 +585,8 @@ def nonzero(fake_mode, func, arg):
             if arg.numel() >= 2:
                 maxval = int(arg.numel())
 
-        _constrain_range_for_size(nnz, max=maxval)
+        from torch.export import constrain_as_size
+        constrain_as_size(nnz)
 
         arg._nonzero_memo = nnz
         arg._nonzero_memo_vc = arg._version
```

there may be fuzz depending on if you patched in https://github.com/pytorch/pytorch/pull/109857 just replace the constrain call with this.

This fails with 

```
  File "/data/users/ezyang/c/pytorch/torch/fx/experimental/proxy_tensor.py", line 574, in __torch_dispatch__                                                                  
    return self.inner_torch_dispatch(func, types, args, kwargs)                                                                                                               
  File "/data/users/ezyang/c/pytorch/torch/fx/experimental/proxy_tensor.py", line 609, in inner_torch_dispatch                                                                
    return proxy_call(self, func, self.pre_dispatch, args, kwargs)                                                                                                            
  File "/data/users/ezyang/c/pytorch/torch/fx/experimental/proxy_tensor.py", line 354, in proxy_call                                                                          
    out = func(*args, **kwargs)                                                                                                                                               
  File "/data/users/ezyang/c/pytorch/torch/_ops.py", line 498, in __call__                                                                                                    
    return self._op(*args, **kwargs or {})                                                                                                                                    
  File "/data/users/ezyang/c/pytorch/torch/utils/_stats.py", line 20, in wrapper                                                                                              
    return fn(*args, **kwargs)                                                                                                                                                
  File "/data/users/ezyang/c/pytorch/torch/_subclasses/fake_tensor.py", line 1300, in __torch_dispatch__                                                                      
    return self.dispatch(func, types, args, kwargs)                                                                                                                           
  File "/data/users/ezyang/c/pytorch/torch/_subclasses/fake_tensor.py", line 1543, in dispatch                                                                                
    op_impl_out = op_impl(self, func, *args, **kwargs)
  File "/data/users/ezyang/c/pytorch/torch/_subclasses/fake_tensor.py", line 587, in nonzero
    torch.sym_constrain_range_for_size(nnz, min=None, max=maxval)
RuntimeError: tried to get Int out of SymInt
```

when you run `python test/test_proxy_tensor.py -k test_make_fx_fake_exhaustive_nonzero_cpu_float32`

I suspect what is going on is you're going to the C++ kernel implementation.

Somehow a simplified repro doesn't trigger though:

```
import torch
from torch.fx.experimental.proxy_tensor import make_fx

def f(x):
    y = x.item()
    print(y)
    torch.sym_constrain_range_for_size(y, min=None, max=None)
    return torch.empty(y)

make_fx(f, tracing_mode="symbolic")(torch.tensor([3]))
```

maybe pydispatcher related?

### Versions

main

cc @msaroufim @bdhirsh @anijain2305 @zou3519 @chauhang @avikchaudhuri @gmagogsfm @zhxchen17 @tugsbayasgalan @angelayi @suo @ydwu4 @wconstab