Repository: pytorch
Issue: Replace pattern fails on incompatible function arguments
Labels: oncall: fx

### 🐛 Describe the bug

replace_pattern fails on incompatible function arguments in Replacement. Expectation is Tensor and it receives a Proxy.

```
import torch
from torch.fx import symbolic_trace, replace_pattern

class LN(torch.nn.Module):
    def __init__(self, h):
        super().__init__()
        self.model = torch.nn.LayerNorm(h)
        
    def forward(self, input):
        out = self.model(input)
        return out
        
class Pattern(torch.nn.Module):
    def __init__(self, h):
        super().__init__()
        self.model = torch.nn.LayerNorm(h)

    def forward(self, input):
        out = self.model(input)
        return out

class Replacement(torch.nn.Module):
    def __init__(self, h):
        super().__init__()
        from apex.contrib.layer_norm.layer_norm import FastLayerNorm
        self.model = FastLayerNorm(h)

    def forward(self, input):
        out = self.model(input)
        return out

def test_ln():
    shape = (4,512,2048)
    in_t = torch.rand(shape, device='cuda', dtype=torch.float16)
    net = LN(2048).cuda().half()
    traced_module = symbolic_trace(net)
    replace_pattern(traced_module, Pattern(2048), Replacement(2048))
    outputs = traced_module(in_t)
    
test_ln()
```

This is the error produced

```
Traceback (most recent call last):
  File "/scratch/mojitos/debug_scripts/temp.py", line 41, in <module>
    test_ln()
  File "/scratch/mojitos/debug_scripts/temp.py", line 37, in test_ln
    replace_pattern(traced_module, Pattern(2048), Replacement(2048))
  File "/opt/pytorch/pytorch/torch/fx/subgraph_rewriter.py", line 196, in replace_pattern
    match_and_replacements = _replace_pattern(gm, pattern, replacement)
  File "/opt/pytorch/pytorch/torch/fx/subgraph_rewriter.py", line 244, in _replace_pattern
    replacement_graph = symbolic_trace(replacement).graph
  File "/opt/pytorch/pytorch/torch/fx/_symbolic_trace.py", line 1070, in symbolic_trace
    graph = tracer.trace(root, concrete_args)
  File "/opt/pytorch/pytorch/torch/fx/_symbolic_trace.py", line 739, in trace
    (self.create_arg(fn(*args)),),
  File "/scratch/mojitos/debug_scripts/temp.py", line 29, in forward
    out = self.model(input)
  File "/opt/pytorch/pytorch/torch/fx/_symbolic_trace.py", line 717, in module_call_wrapper
    return self.call_module(mod, forward, args, kwargs)
  File "/opt/pytorch/pytorch/torch/fx/_symbolic_trace.py", line 434, in call_module
    return forward(*args, **kwargs)
  File "/opt/pytorch/pytorch/torch/fx/_symbolic_trace.py", line 710, in forward
    return _orig_module_call(mod, *args, **kwargs)
  File "/opt/pytorch/pytorch/torch/nn/modules/module.py", line 1488, in _call_impl
    return forward_call(*args, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/apex/contrib/layer_norm/layer_norm.py", line 53, in forward
    return _fast_layer_norm(x, self.weight, self.bias, self.epsilon)
  File "/usr/local/lib/python3.10/site-packages/apex/contrib/layer_norm/layer_norm.py", line 37, in _fast_layer_norm
    return FastLayerNormFN.apply(*args)
  File "/opt/pytorch/pytorch/torch/autograd/function.py", line 453, in apply
    return super().apply(*args, **kwargs)
  File "/usr/local/lib/python3.10/site-packages/apex/contrib/layer_norm/layer_norm.py", line 16, in forward
    ymat, mu, rsigma = fast_layer_norm.ln_fwd(xmat, gamma, beta, epsilon)
TypeError: ln_fwd(): incompatible function arguments. The following argument types are supported:
    1. (arg0: torch.Tensor, arg1: torch.Tensor, arg2: torch.Tensor, arg3: float) -> List[torch.Tensor]

Invoked with: Proxy(view), Proxy(contiguous_1), Proxy(contiguous_2), 1e-05
```

### Versions

```
Collecting environment information...
PyTorch version: 2.0.0a0+git5e0d345
Is debug build: False
CUDA used to build PyTorch: 12.0
ROCM used to build PyTorch: N/A

OS: Ubuntu 20.04.5 LTS (x86_64)
GCC version: (Ubuntu 9.4.0-1ubuntu1~20.04.1) 9.4.0
Clang version: Could not collect
CMake version: version 3.25.0
Libc version: glibc-2.31

Python version: 3.10.9 (main, Jan 17 2023, 00:37:25) [GCC 9.4.0] (64-bit runtime)
Python platform: Linux-4.15.0-142-generic-x86_64-with-glibc2.31
Is CUDA available: True
CUDA runtime version: 12.0.140
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA A100 80GB PCIe
Nvidia driver version: 525.60.13
cuDNN version: Probably one of the following:
/usr/lib/x86_64-linux-gnu/libcudnn.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_infer.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_adv_train.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_infer.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_cnn_train.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_infer.so.8.7.0
/usr/lib/x86_64-linux-gnu/libcudnn_ops_train.so.8.7.0
HIP runtime version: N/A
MIOpen runtime version: N/A
Is XNNPACK available: True

Versions of relevant libraries:
[pip3] numpy==1.23.5
[pip3] torch==2.0.0a0+git5e0d345
[pip3] torchvision==0.15.0a0+035d99f
[conda] Could not collect
```

cc @ezyang @SherlockNoMad @EikanWang @jgong5 @wenzhe-nrv