Hi, when I am trying to compile torch scatter C++ API, I did the cmake as suggested in the readme. However, when the cmake trying to find python3, it keeps finding an old version of python. I have tried numerous ways (like add set() or inlcude_directories or -D tags when doing the cmake line), but it either keeps finding the old python3.8 or bugged out by saying it can not find a python3.10, when python3.10 is the env I configed for my applications. Would you mind providing some examples on how to modify cmakes or other suggestions so that I can force the cmake (particularly the find_package() function) to use my python3.10 to compile scatter C++ API?

Thank you so much for the help.

Error during compilation of extension

Reproduce:

1. Install Pytorch 2.0 python3 -m pip install numpy --pre torch --force-reinstall --extra-index-url https://download.pytorch.org/whl/nightly/cu117
2. Try to build using python3 -m pip install torch-scatter

I'm trying to implement a scatter_sum operation via the c++ api.

I'm calling the function as follows:

results = scatter_sum(source_nodes, target_index_tensor, dim, torch::nullopt, torch::nullopt);

I have verified that both tensors are on cuda:0 via these lines:

std::cout << source_nodes.device() << std::endl;
std::cout << target_index_tensor.device() << std::endl;

The program simply fails when I used 'kCUDA' as the device, but when I use 'kCPU' as the device, it works. I have verified that the normal torch functions (linear, relu) work on the kCUDA device, so only this scatter_sum function does not go through. What could be the cause of the program failing? I simply get 'core dumped', but because it works on CPU, it's not so clear to me what could be wrong.

Some information about the system: Python 3.8 CUDA 10.2 PyTorch 1.10

Overview

Hi 👋🏼 ,

I would just like to start by saying, thank you for creating and maintaining this amazing library.

When attempting to use functorch with pytorch-geometric I encountered the follwoing error related to scatter_add. Please let me know if I can provide anymore information or help out in anyway.

Thank you, Matt

Code

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)

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.

When I run the following code:

a = torch.arange(10)
indptr = torch.tensor([0]) # invalid ptr
segment_csr(a, indptr)

Python crashes on OSX with the following message:

I'm on version 2.0.9 of torch_scatter. I think segment_csr should check for bad input like this.