Repository: pytorch
Issue: Dtype changes while going from FX graph -> Torchscript
Labels: triaged,FX-TorchScript Compatibility,module: fx

### 🐛 Describe the bug

Python:
```
rand_idx = torch.randint(sy*sx, size=(hsy, wsx, 1), device=metric.device, dtype=torch.int64)
idx_buffer_view = torch.zeros(hsy, wsx, sy*sx, device=metric.device, dtype=torch.int64)
idx_buffer_view.scatter_(dim=2, index=rand_idx, src=-torch.ones_like(rand_idx, dtype=rand_idx.dtype))
```

The `FX graph` :
```
%randint : [#users=2] = call_function[target=torch.ops.aten.randint.default](args = (4, [32, 32, 1]), kwargs = {device: cpu, pin_memory: False})
%zeros : [#users=1] = call_function[target=torch.ops.aten.zeros.default](args = ([32, 32, 4],), kwargs = {dtype: torch.int64, device: cpu, pin_memory: False})
%ones_like : [#users=1] = call_function[target=torch.ops.aten.ones_like.default](args = (%randint,), kwargs = {dtype: torch.int64, pin_memory: False})
%neg : [#users=1] = call_function[target=torch.ops.aten.neg.default](args = (%ones_like,), kwargs = {})
%scatter_ : [#users=1] = call_function[target=torch.ops.aten.scatter_.src](args = (%zeros, 2, %randint, %neg), kwargs = {})
```    
The `Torchscript` IR :
```
%836 = torch.aten.randint %int4, %835, %int4, %none_3, %cpu, %false : !torch.int, !torch.list<int>, !torch.int, !torch.none, !torch.Device, !torch.bool -> !torch.tensor
%838 = torch.aten.zeros %837, %int5, %none_3, %cpu, %false : !torch.list<int>, !torch.int, !torch.none, !torch.Device, !torch.bool -> !torch.tensor
%839 = torch.aten.ones_like %836, %int4, %none_3, %none_3, %false, %none_3 : !torch.tensor, !torch.int, !torch.none, !torch.none, !torch.bool, !torch.none -> !torch.tensor
%840 = torch.aten.neg %839 : !torch.tensor -> !torch.tensor
%841 = torch.aten.scatter_.src %838, %int2, %836, %840 : !torch.tensor, !torch.int, !torch.tensor, !torch.tensor -> !torch.tensor
```

`torch.aten.randint ` : 3rd argument is `dtype`, in this case it's `%int4` (int64)
`torch.aten.zeros` : 2nd argument is `dtype`, in this case it's `%int5`. (half)
`torch.aten.ones_like` : 2nd argument is `dtype`, in this case it's `%int4`. (int64)

The reason behind `torch.aten.zeros` being set to have dtype as`fp16` despite having `int64` in the Python code is because when an FX graph is converted to TorchScript and imported into Torch-MLIR, a Python representation of the graph is included as a string parameter in the MLIR module. All the `torch.ops.aten.zeros` in this Python representation are being set to have `dtyp = torch.float16` (that's a bug!).

You can observe that in the [Torchscript](https://drive.google.com/file/d/167bzYzDKfv6G1WR5Gnv0vAHDSUoGIVrN/view?usp=sharing) IR file.
The following is what you'd observe :-
```
zeros = torch.ops.aten.zeros([32, 32, 4], dtype = torch.float16, device = device(type='cpu'), pin_memory = False)
zeros_1 = torch.ops.aten.zeros([2, 4096, 320], dtype = torch.float16, device = device(type='cpu'), pin_memory = False)
zeros_2 = torch.ops.aten.zeros([32, 32, 4], dtype = torch.float16, device = device(type='cpu'), pin_memory = False)
zeros_3 = torch.ops.aten.zeros([2, 4096, 320], dtype = torch.float16, device = device(type='cpu'), pin_memory = False)
zeros_4 = torch.ops.aten.zeros([32, 32, 4], dtype = torch.float16, device = device(type='cpu'), pin_memory = False)
zeros_5 = torch.ops.aten.zeros([2, 4096, 320], dtype = torch.float16, device = device(type='cpu'), pin_memory = False)
zeros_6 = torch.ops.aten.zeros([32, 32, 4], dtype = torch.float16, device = device(type='cpu'), pin_memory = False)
zeros_7 = torch.ops.aten.zeros([2, 4096, 320], dtype = torch.float16, device = device(type='cpu'), pin_memory = False)
zeros_8 = torch.ops.aten.zeros([32, 32, 4], dtype = torch.float16, device = device(type='cpu'), pin_memory = False)
zeros_9 = torch.ops.aten.zeros([2, 4096, 320], dtype = torch.float16, device = device(type='cpu'), pin_memory = False)
```

Also, here is the corresponding [FX graph](https://drive.google.com/file/d/1KQrOEtDzoUjf-H_oee_Un30QME05Lwfk/view?usp=sharing) file just for your reference.

### Versions

```
Collecting environment information...
PyTorch version: 2.1.0.dev20230403+cu118
Is debug build: False
CUDA used to build PyTorch: 11.8
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04.1 LTS (x86_64)
GCC version: (Ubuntu 11.3.0-1ubuntu1~22.04) 11.3.0
Clang version: 14.0.0-1ubuntu1
CMake version: version 3.26.3
Libc version: glibc-2.35

Python version: 3.11.2 (main, Feb  8 2023, 14:49:25) [GCC 11.3.0] (64-bit runtime)
Python platform: Linux-5.15.0-1030-gcp-x86_64-with-glibc2.35
Is CUDA available: True
CUDA runtime version: 12.0.76
CUDA_MODULE_LOADING set to: LAZY
GPU models and configuration: GPU 0: NVIDIA A100-SXM4-40GB
Nvidia driver version: 525.85.12
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
/usr/local/cuda-12.0/targets/x86_64-linux/lib/libcudnn.so.8.7.0
/usr/local/cuda-12.0/targets/x86_64-linux/lib/libcudnn_adv_infer.so.8.7.0
/usr/local/cuda-12.0/targets/x86_64-linux/lib/libcudnn_adv_train.so.8.7.0
/usr/local/cuda-12.0/targets/x86_64-linux/lib/libcudnn_cnn_infer.so.8.7.0
/usr/local/cuda-12.0/targets/x86_64-linux/lib/libcudnn_cnn_train.so.8.7.0
/usr/local/cuda-12.0/targets/x86_64-linux/lib/libcudnn_ops_infer.so.8.7.0
/usr/local/cuda-12.0/targets/x86_64-linux/lib/libcudnn_ops_train.so.8.7.0
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
Address sizes:                   46 bits physical, 48 bits virtual
Byte Order:                      Little Endian
CPU(s):                          12
On-line CPU(s) list:             0-11
Vendor ID:                       GenuineIntel
Model name:                      Intel(R) Xeon(R) CPU @ 2.20GHz
CPU family:                      6
Model:                           85
Thread(s) per core:              2
Core(s) per socket:              6
Socket(s):                       1
Stepping:                        7
BogoMIPS:                        4400.39
Flags:                           fpu vme de pse tsc msr pae mce cx8 apic sep mtrr pge mca cmov pat pse36 clflush mmx fxsr sse sse2 ss ht syscall nx pdpe1gb rdtscp lm constant_tsc rep_good nopl xtopology nonstop_tsc cpuid tsc_known_freq pni pclmulqdq ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single ssbd ibrs ibpb stibp ibrs_enhanced fsgsbase tsc_adjust bmi1 hle avx2 smep bmi2 erms invpcid rtm mpx avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xgetbv1 xsaves arat avx512_vnni md_clear arch_capabilities
Hypervisor vendor:               KVM
Virtualization type:             full
L1d cache:                       192 KiB (6 instances)
L1i cache:                       192 KiB (6 instances)
L2 cache:                        6 MiB (6 instances)
L3 cache:                        38.5 MiB (1 instance)
NUMA node(s):                    1
NUMA node0 CPU(s):               0-11
Vulnerability Itlb multihit:     Not affected
Vulnerability L1tf:              Not affected
Vulnerability Mds:               Mitigation; Clear CPU buffers; SMT Host state unknown
Vulnerability Meltdown:          Not affected
Vulnerability Mmio stale data:   Vulnerable: Clear CPU buffers attempted, no microcode; SMT Host state unknown
Vulnerability Retbleed:          Mitigation; Enhanced IBRS
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
Vulnerability Srbds:             Not affected
Vulnerability Tsx async abort:   Mitigation; Clear CPU buffers; SMT Host state unknown

Versions of relevant libraries:
[pip3] mypy-extensions==1.0.0
[pip3] numpy==1.24.2
[pip3] pytorch-lightning==2.0.1.post0
[pip3] pytorch-triton==2.1.0+46672772b4
[pip3] torch==2.1.0.dev20230403+cu118
[pip3] torch-mlir==20230411.805
[pip3] torchmetrics==0.11.4
[pip3] torchvision==0.16.0.dev20230403+cu118
[pip3] triton==2.0.0
[conda] Could not collect
```

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