Repository: pytorch
Issue: [JIT][foreach_add] Arguments for call are not valid
Labels: oncall: jit

### 🐛 Describe the bug

Observe an error when try to use torch.jit.script on _foreach_add in following script:
```
import torch


def func_cpu():
  input = [torch.ones(2,2)]
  scalars = [0.5]
  return torch._foreach_add(input, scalars)[0]



print(func_cpu()) # Works as use aten::_foreach_add.ScalarList
print(torch.jit.script(func_cpu)) # Fails to create a jit function.
```
the error is following:
```
tensor([[1.5000, 1.5000],
        [1.5000, 1.5000]])
Traceback (most recent call last):
  File "/home/pswider/foreach_add_repro.py", line 12, in <module>
    print(torch.jit.script(func_cpu))
  File "/home/pswider/.local/lib/python3.10/site-packages/torch/jit/_script.py", line 1432, in script
    return _script_impl(
  File "/home/pswider/.local/lib/python3.10/site-packages/torch/jit/_script.py", line 1204, in _script_impl
    fn = torch._C._jit_script_compile(
RuntimeError: 
Arguments for call are not valid.
The following variants are available:
  
  aten::_foreach_add.Scalar(Tensor[] self, Scalar scalar) -> Tensor[]:
  Expected a value of type 'number' for argument 'scalar' but instead found type 'List[float]'.
  
  aten::_foreach_add.Scalar_out(Tensor[] self, Scalar scalar, *, Tensor(a!)[] out) -> ():
  Expected a value of type 'number' for argument 'scalar' but instead found type 'List[float]'.
  
  aten::_foreach_add.List(Tensor[] self, Tensor[] other, *, Scalar alpha=1) -> Tensor[]:
  Expected a value of type 'List[Tensor]' for argument 'other' but instead found type 'List[float]'.
  
  aten::_foreach_add.List_out(Tensor[] self, Tensor[] other, *, Scalar alpha=1, Tensor(a!)[] out) -> ():
  Expected a value of type 'List[Tensor]' for argument 'other' but instead found type 'List[float]'.
  
  aten::_foreach_add.ScalarList(Tensor[] self, Scalar[] scalars) -> Tensor[]:
  Expected a value of type 'List[number]' for argument 'scalars' but instead found type 'List[float]'.
  
  aten::_foreach_add.ScalarList_out(Tensor[] self, Scalar[] scalars, *, Tensor(a!)[] out) -> ():
  Expected a value of type 'List[number]' for argument 'scalars' but instead found type 'List[float]'.
  
  aten::_foreach_add.Tensor(Tensor[] self, Tensor other, *, Scalar alpha=1) -> Tensor[]:
  Expected a value of type 'Tensor' for argument 'other' but instead found type 'List[float]'.
  
  aten::_foreach_add.Tensor_out(Tensor[] self, Tensor other, *, Scalar alpha=1, Tensor(a!)[] out) -> ():
  Expected a value of type 'Tensor' for argument 'other' but instead found type 'List[float]'.

The original call is:
  File "/home/pswider/foreach_add_repro.py", line 7
  input = [torch.ones(2,2)]
  scalars = [0.5]
  return torch._foreach_add(input, scalars)[0]
         ~~~~~~~~~~~~~~~~~~ <--- HERE
```

TorchScript seems to treat List[float] as it isn't List[number] which cause RuntimeError

### Versions

Collecting environment information...
PyTorch version: 2.4.0+cu121
Is debug build: False
CUDA used to build PyTorch: 12.1
ROCM used to build PyTorch: N/A

OS: Ubuntu 22.04 LTS (x86_64)
GCC version: (Ubuntu 11.4.0-1ubuntu1~22.04) 11.4.0
Clang version: 14.0.0-1ubuntu1.1
CMake version: version 3.27.9
Libc version: glibc-2.35

Python version: 3.10.12 (main, Jul 29 2024, 16:56:48) [GCC 11.4.0] (64-bit runtime)
Python platform: Linux-5.15.0-117-generic-x86_64-with-glibc2.35
Is CUDA available: False
CUDA runtime version: No CUDA
CUDA_MODULE_LOADING set to: N/A
GPU models and configuration: No CUDA
Nvidia driver version: No CUDA
cuDNN version: No CUDA
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:                        43 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) Gold 6132 CPU @ 2.60GHz
CPU family:                           6
Model:                                85
Thread(s) per core:                   1
Core(s) per socket:                   6
Socket(s):                            2
Stepping:                             0
BogoMIPS:                             5187.81
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 arch_perfmon nopl xtopology tsc_reliable nonstop_tsc cpuid tsc_known_freq pni pclmulqdq vmx ssse3 fma cx16 pcid sse4_1 sse4_2 x2apic movbe popcnt tsc_deadline_timer aes xsave avx f16c rdrand hypervisor lahf_lm abm 3dnowprefetch invpcid_single pti ssbd ibrs ibpb stibp tpr_shadow vnmi ept vpid ept_ad fsgsbase tsc_adjust bmi1 avx2 smep bmi2 invpcid avx512f avx512dq rdseed adx smap clflushopt clwb avx512cd avx512bw avx512vl xsaveopt xsavec xsaves arat pku ospke md_clear flush_l1d arch_capabilities
Virtualization:                       VT-x
Hypervisor vendor:                    VMware
Virtualization type:                  full
L1d cache:                            384 KiB (12 instances)
L1i cache:                            384 KiB (12 instances)
L2 cache:                             12 MiB (12 instances)
L3 cache:                             38.5 MiB (2 instances)
NUMA node(s):                         1
NUMA node0 CPU(s):                    0-11
Vulnerability Gather data sampling:   Unknown: Dependent on hypervisor status
Vulnerability Itlb multihit:          KVM: Mitigation: VMX disabled
Vulnerability L1tf:                   Mitigation; PTE Inversion; VMX flush not necessary, SMT disabled
Vulnerability Mds:                    Mitigation; Clear CPU buffers; SMT Host state unknown
Vulnerability Meltdown:               Mitigation; PTI
Vulnerability Mmio stale data:        Mitigation; Clear CPU buffers; SMT Host state unknown
Vulnerability Reg file data sampling: Not affected
Vulnerability Retbleed:               Mitigation; IBRS
Vulnerability Spec rstack overflow:   Not affected
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; IBRS; IBPB conditional; STIBP disabled; RSB filling; PBRSB-eIBRS Not affected; BHI SW loop, KVM SW loop
Vulnerability Srbds:                  Not affected
Vulnerability Tsx async abort:        Not affected

Versions of relevant libraries:
[pip3] flake8==6.1.0
[pip3] numpy==1.23.5
[pip3] torch==2.4.0
[pip3] triton==3.0.0
[conda] Could not collect

cc @EikanWang @jgong5 @wenzhe-nrv @sanchitintel