Repository: pytorch
Issue: Dropout traces poorly with AotAutograd/make_fx
Labels: triaged,oncall: pt2,module: aotdispatch,module: pt2-dispatcher

## Issue 1: probability can't be a tensor
```py
from torch.fx.experimental.proxy_tensor import make_fx
import torch

print(make_fx(
    lambda x,p: torch.nn.functional.dropout(x, p),
    decomposition_table={},
    tracing_mode="fake")(torch.randn(10), torch.tensor(0.5)))
```
output:
```py
Traceback (most recent call last):
  File "repro.py", line 4, in <module>
    print(make_fx(
  File "/home/jansel/pytorch/torch/fx/experimental/proxy_tensor.py", line 756, in wrapped
    t = dispatch_trace(wrap_key(func, args, fx_tracer, pre_autograd), tracer=fx_tracer, concrete_args=tuple(phs))
  File "/home/jansel/pytorch/torch/fx/experimental/proxy_tensor.py", line 462, in dispatch_trace
    graph = tracer.trace(root, concrete_args)
  File "/home/jansel/pytorch/torch/fx/_symbolic_trace.py", line 778, in trace
    (self.create_arg(fn(*args)),),
  File "/home/jansel/pytorch/torch/fx/experimental/proxy_tensor.py", line 479, in wrapped
    out = f(*tensors)
  File "repro.py", line 5, in <lambda>
    lambda x, p: torch.nn.functional.dropout(x, p),
  File "/home/jansel/pytorch/torch/nn/functional.py", line 1250, in dropout
    if p < 0.0 or p > 1.0:
  File "/home/jansel/pytorch/torch/utils/_stats.py", line 20, in wrapper
    return fn(*args, **kwargs)
  File "/home/jansel/pytorch/torch/fx/experimental/proxy_tensor.py", line 528, in __torch_dispatch__
    return self.inner_torch_dispatch(func, types, args, kwargs)
  File "/home/jansel/pytorch/torch/fx/experimental/proxy_tensor.py", line 553, in inner_torch_dispatch
    return proxy_call(self, func, self.pre_autograd, args, kwargs)
  File "/home/jansel/pytorch/torch/fx/experimental/proxy_tensor.py", line 364, in proxy_call
    out = func(*args, **kwargs)
  File "/home/jansel/pytorch/torch/_ops.py", line 394, in __call__
    return self._op(*args, **kwargs or {})
  File "/home/jansel/pytorch/torch/utils/_stats.py", line 20, in wrapper
    return fn(*args, **kwargs)
  File "/home/jansel/pytorch/torch/_subclasses/fake_tensor.py", line 1062, in __torch_dispatch__
    return self.dispatch(func, types, args, kwargs)
  File "/home/jansel/pytorch/torch/_subclasses/fake_tensor.py", line 1244, in dispatch
    op_impl_out = op_impl(self, func, *args, **kwargs)
  File "/home/jansel/pytorch/torch/_subclasses/fake_tensor.py", line 415, in local_scalar_dense
    raise DataDependentOutputException(func)
torch._subclasses.fake_tensor.DataDependentOutputException: aten._local_scalar_dense.default
```


You can also hit this issue with `mode="real"`
```py
print(make_fx(
    lambda x, p: torch.nn.functional.dropout(x, p),
    decomposition_table={},
    tracing_mode="real")(torch.randn(10, device="cuda"), torch.tensor(0.235, device="cuda")))
```
outputs:
```py
Traceback (most recent call last):
  File "repro.py", line 4, in <module>
    print(make_fx(
  File "/home/jansel/pytorch/torch/fx/experimental/proxy_tensor.py", line 756, in wrapped
    t = dispatch_trace(wrap_key(func, args, fx_tracer, pre_autograd), tracer=fx_tracer, concrete_args=tuple(phs))
  File "/home/jansel/pytorch/torch/fx/experimental/proxy_tensor.py", line 462, in dispatch_trace
    graph = tracer.trace(root, concrete_args)
  File "/home/jansel/pytorch/torch/fx/_symbolic_trace.py", line 778, in trace
    (self.create_arg(fn(*args)),),
  File "/home/jansel/pytorch/torch/fx/experimental/proxy_tensor.py", line 479, in wrapped
    out = f(*tensors)
  File "repro.py", line 5, in <lambda>
    lambda x, p: torch.nn.functional.dropout(x, p),
  File "/home/jansel/pytorch/torch/nn/functional.py", line 1250, in dropout
    if p < 0.0 or p > 1.0:
  File "/home/jansel/pytorch/torch/utils/_stats.py", line 20, in wrapper
    return fn(*args, **kwargs)
  File "/home/jansel/pytorch/torch/fx/experimental/proxy_tensor.py", line 528, in __torch_dispatch__
    return self.inner_torch_dispatch(func, types, args, kwargs)
  File "/home/jansel/pytorch/torch/fx/experimental/proxy_tensor.py", line 553, in inner_torch_dispatch
    return proxy_call(self, func, self.pre_autograd, args, kwargs)
  File "/home/jansel/pytorch/torch/fx/experimental/proxy_tensor.py", line 294, in proxy_call
    raise RuntimeError(
RuntimeError: It appears that you're trying to get value out of a tracing tensor with aten._local_scalar_dense.default - erroring out! It's likely that this is caused by data-dependent control flow or similar.  It may be possible to trace this with dynamic shapes; try setting tracing_mode='symbolic' in your make_fx call.
```



## Issue 2: it gets decomposed unconditionally on CPU

```py
print(make_fx(
    lambda x, p: torch.nn.functional.dropout(x, p),
    decomposition_table={},
    tracing_mode="real")(torch.randn(10), 0.5))
```
results in:
```py
def forward(self, x_1, p_1):
    empty_like = torch.ops.aten.empty_like.default(x_1, memory_format = torch.contiguous_format)
    bernoulli_ = torch.ops.aten.bernoulli_.float(empty_like);  empty_like = None
    div_ = torch.ops.aten.div_.Scalar(bernoulli_, 0.5);  bernoulli_ = None
    mul = torch.ops.aten.mul.Tensor(x_1, div_);  x_1 = div_ = None
    return mul
```

I set `decomposition_table={}` why am I getting `bernoulli_`???

## Issue 3: GPU does something completely different

```py
print(make_fx(
    lambda x, p: torch.nn.functional.dropout(x, p),
    decomposition_table={},
    tracing_mode="real")(torch.randn(10, device="cuda"), 0.5))
```
outputs:
```py
def forward(self, x_1, p_1):
    native_dropout = torch.ops.aten.native_dropout.default(x_1, 0.5, True);  x_1 = None
    getitem = native_dropout[0]
    getitem_1 = native_dropout[1];  native_dropout = None
    return getitem
```
This is almost right, except 0.5 got burned in rather than passed as an arg.



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