Repository: pytorch
Issue: aot_function cannot handle index_add
Labels: triaged,module: functorch

### 🐛 Describe the bug

The following code fails:
```python
import torch
from torch import Tensor
from functorch.compile import aot_function, ts_compile


def foo(positions: Tensor, indexes: Tensor, source: Tensor) -> Tensor:
    n_node = positions.shape[0]
    result = torch.zeros(n_node, 2, dtype=torch.float32, device=positions.device)
    node_features = result.index_add(
        0,
        indexes,
        source,
    )
    return node_features


n_node = 100
n_edge = 1000
positions = torch.randn(n_node, 3, dtype=torch.float32)
indexes = torch.randint(0, n_node, (n_edge,), dtype=torch.long)
source = torch.randn(n_edge, 2, dtype=torch.float32)
foo(positions, indexes, source)
foo = aot_function(foo, fw_compiler=ts_compile) 
foo(positions, indexes, source) # Note: the bug is triggered here
```
printing this error:
```shell
;; This buffer is for text that is not saved, and for Lisp evaluation.
;; To create a file, visit it with <open> and enter text in its buffer.

python functorch_test.py 
/home/raul/mambaforge/envs/all/lib/python3.11/site-packages/torch/jit/_check.py:177: UserWarning: The TorchScript type system doesn't support instance-level annotations on empty non-base types in `__init__`. Instead, either 1) use a type annotation in the class body, or 2) wrap the type in `torch.jit.Attribute`.
  warnings.warn(
Traceback (most recent call last):
  File "/home/raul/torchmd-net/tests/functorch_test.py", line 20, in <module>
    foo(positions, indexes, source)
  File "/home/raul/mambaforge/envs/all/lib/python3.11/site-packages/torch/_functorch/aot_autograd.py", line 778, in returned_function
    out = cached_fn(flat_args)
          ^^^^^^^^^^^^^^^^^^^^
  File "/home/raul/mambaforge/envs/all/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 215, in runtime_wrapper
    all_outs = call_func_at_runtime_with_args(
               ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/raul/mambaforge/envs/all/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/utils.py", line 118, in call_func_at_runtime_with_args
    out = normalize_as_list(f(args))
                            ^^^^^^^
  File "/home/raul/mambaforge/envs/all/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/runtime_wrappers.py", line 450, in wrapper
    return compiled_fn(runtime_args)
           ^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/raul/mambaforge/envs/all/lib/python3.11/site-packages/torch/_functorch/_aot_autograd/utils.py", line 92, in g
    return f(*args)
           ^^^^^^^^
  File "/home/raul/mambaforge/envs/all/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1545, in _wrapped_call_impl
    return self._call_impl(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
  File "/home/raul/mambaforge/envs/all/lib/python3.11/site-packages/torch/nn/modules/module.py", line 1554, in _call_impl
    return forward_call(*args, **kwargs)
           ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
RuntimeError: The following operation failed in the TorchScript interpreter.
Traceback of TorchScript (most recent call last):
  File "<eval_with_key>.3 from /home/raul/mambaforge/envs/all/lib/python3.11/site-packages/torch/fx/experimental/proxy_tensor.py:652 in wrapped", line 6, in forward
def forward(self, arg0_1, arg1_1, arg2_1):
    zeros = torch.ops.aten.zeros([100, 2], dtype = torch.float32, device = 'cuda', pin_memory = False)
    index_add = torch.ops.aten.index_add(zeros, 0, arg1_1, arg2_1);  zeros = arg1_1 = arg2_1 = None
                ~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    return (index_add,)
RuntimeError: 
```
Note that `foo` can be TorchScript'ed just fine.

### Versions

I am using a recent nightly
```shell
Versions of relevant libraries:
[pip3] numpy==1.26.4
[pip3] torch==2.4.0.dev20240531
[pip3] triton==3.0.0
[conda] blas                      2.116                       mkl    conda-forge
[conda] blas-devel                3.9.0            16_linux64_mkl    conda-forge
[conda] libblas                   3.9.0            16_linux64_mkl    conda-forge
[conda] libcblas                  3.9.0            16_linux64_mkl    conda-forge
[conda] liblapack                 3.9.0            16_linux64_mkl    conda-forge
[conda] liblapacke                3.9.0            16_linux64_mkl    conda-forge
[conda] mkl                       2022.1.0           h84fe81f_915    conda-forge
[conda] mkl-devel                 2022.1.0           ha770c72_916    conda-forge
[conda] mkl-include               2022.1.0           h84fe81f_915    conda-forge
[conda] numpy                     1.26.4                   pypi_0    pypi
[conda] pytorch                   2.4.0.dev20240531 py3.11_cuda12.4_cudnn8.9.2_0    pytorch-nightly
[conda] pytorch-cuda              12.4                 hc786d27_6    pytorch-nightly
[conda] pytorch-mutex             1.0                        cuda    pytorch-nightly
[conda] torchtriton               3.0.0+45fff310c8           py311    pytorch-nightly
```

cc @zou3519 @Chillee @samdow @kshitij12345 @janeyx99