Repository: pytorch
Issue: functorch aten::scatter_add_  not implemented
Labels: triaged,module: functorch

### 🐛 Describe the bug

Hi 👋🏼 ,

Thank you for maintaining this amazing project.

When attempting to use `functorch` with `torch-geometric` I encounter the following error. Please let me know if there is any additional information I can provide or anything I can do to help.

Thank you,
Matt
```python
from functorch import combine_state_for_ensemble, vmap
from torch import nn
from torch_geometric.nn import GCNConv
from torch_geometric.data import Data
import torch

NUM_MODELS = 10
INPUT_SIZE = 8
NUM_NODES, NUM_EDGES = 4, 8

# create a model
class Model(nn.Module):
    def __init__(self, input_size: int) -> None:
        super().__init__()
        self.conv1 = GCNConv(input_size, 2, add_self_loops=False).jittable()

    
    def forward(self, x: torch.Tensor, edge_index: torch.Tensor) -> torch.Tensor:
        return self.conv1(x, edge_index)

# create the data
xs = torch.randn(NUM_MODELS, NUM_NODES, INPUT_SIZE, dtype=torch.float)
edge_indices = torch.randint(0, 3, (NUM_MODELS, 2, NUM_EDGES), dtype=torch.long)

# create functional models
models = [Model(INPUT_SIZE) for _ in range(NUM_MODELS)]
fmodel, params, buffers = combine_state_for_ensemble(models)

# complete a forward pass with the data
res = vmap(fmodel)(params, buffers, xs, edge_indices)
```
Produces the error...
```
(.venv) matthewlemay@Matthews-MacBook-Air ~/G/torch-func [0|1]> python3 run.py
/Users/matthewlemay/Github/torch-func/.venv/lib/python3.10/site-packages/torch_scatter/scatter.py:21: UserWarning: There is a performance drop because we have not yet implemented the batching rule for aten::scatter_add_. Please file us an issue on GitHub so that we can prioritize its implementation. (Triggered internally at  /Users/runner/work/functorch/functorch/functorch/csrc/BatchedFallback.cpp:85.)
  return out.scatter_add_(dim, index, src)
Traceback (most recent call last):
  File "/Users/matthewlemay/Github/torch-func/run.py", line 30, in <module>
    res = vmap(fmodel)(params, buffers, xs, edge_indices)
  File "/Users/matthewlemay/Github/torch-func/.venv/lib/python3.10/site-packages/functorch/_src/vmap.py", line 365, in wrapped
    batched_outputs = func(*batched_inputs, **kwargs)
  File "/Users/matthewlemay/Github/torch-func/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Users/matthewlemay/Github/torch-func/.venv/lib/python3.10/site-packages/functorch/_src/make_functional.py", line 282, in forward
    return self.stateless_model(*args, **kwargs)
  File "/Users/matthewlemay/Github/torch-func/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/Users/matthewlemay/Github/torch-func/run.py", line 19, in forward
    return self.conv1(x, edge_index)
  File "/Users/matthewlemay/Github/torch-func/.venv/lib/python3.10/site-packages/torch/nn/modules/module.py", line 1130, in _call_impl
    return forward_call(*input, **kwargs)
  File "/var/folders/hh/vh54hqf544n7qf9vxn1lt8_00000gn/T/matthewlemay_pyg/tmp97j0p1uv.py", line 219, in forward
    edge_index, edge_weight = gcn_norm(  # yapf: disable
  File "/Users/matthewlemay/Github/torch-func/.venv/lib/python3.10/site-packages/torch_geometric/nn/conv/gcn_conv.py", line 67, in gcn_norm
    deg = scatter_add(edge_weight, idx, dim=0, dim_size=num_nodes)
  File "/Users/matthewlemay/Github/torch-func/.venv/lib/python3.10/site-packages/torch_scatter/scatter.py", line 29, in scatter_add
    return scatter_sum(src, index, dim, out, dim_size)
  File "/Users/matthewlemay/Github/torch-func/.venv/lib/python3.10/site-packages/torch_scatter/scatter.py", line 21, in scatter_sum
    return out.scatter_add_(dim, index, src)
RuntimeError: vmap: aten::scatter_add_(self, *extra_args) is not possible because there exists a Tensor `other` in extra_args that has more elements than `self`. This happened due to `other` being vmapped over but `self` not being vmapped over at level 1. Please try to use out-of-place operators instead of aten::scatter_add_. If said operator is being called inside the PyTorch framework, please file a bug report instead.
```

### Versions

```
PyTorch version: 1.12.1
Is debug build: False
CUDA used to build PyTorch: None
ROCM used to build PyTorch: N/A

OS: macOS 12.6 (x86_64)
GCC version: Could not collect
Clang version: 14.0.0 (clang-1400.0.29.102)
CMake version: version 3.24.2
Libc version: N/A

Python version: 3.10.6 (main, Aug 30 2022, 05:12:36) [Clang 13.1.6 (clang-1316.0.21.2.5)] (64-bit runtime)
Python platform: macOS-12.6-x86_64-i386-64bit
Is CUDA available: False
CUDA runtime version: No CUDA
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

Versions of relevant libraries:
[pip3] functorch==0.2.1
[pip3] numpy==1.23.3
[pip3] torch==1.12.1
[pip3] torch-cluster==1.6.0
[pip3] torch-geometric==2.1.0.post1
[pip3] torch-scatter==2.0.9
[pip3] torch-sparse==0.6.15
[pip3] torch-spline-conv==1.2.1
[conda] Could not collect
```

cc @zou3519 @Chillee @samdow @soumith