torch::jit::load error

[qiuqiouba]

🐛 Describe the bug

I have done script my model in python, and want load it in c++. It can be confirmed that torch_scatcher and torch_sparse have been successfully compiled.
I can run the following example and get the correct result.

#include <torch/script.h>
#include <torch/torch.h>
#include <pytorch_scatter/scatter.h>
#include <pytorch_sparse/sparse.h>
#include <iostream>
int main()
{
    torch::Tensor src = torch::tensor({ 0.5, 0.4, 0.1, 0.6 });
    torch::Tensor index = torch::tensor({ 0, 0, 1, 1 });
    std::cout << src << std::endl;
    std::cout << index << std::endl;
    std::cout << torch::cuda::cudnn_is_available() << std::endl;
    std::cout << torch::cuda::is_available() << std::endl;
    std::cout << scatter_sum(src, index, 0, torch::nullopt, torch::nullopt) << std::endl;
    torch::Tensor tensor = torch::tensor({ 0, 0, 0, 0, 1, 1, 1});
    std::cout << tensor << std::endl;
    std::cout << ind2ptr(tensor, 2) << std::endl;
}

output:
 0.5000
 0.4000
 0.1000
 0.6000
[ CPUFloatType{4} ]
 0
 0
 1
 1
[ CPULongType{4} ]
1
1
 0.9000
 0.7000
[ CPUFloatType{2} ]
 0
 0
 0
 0
 1
 1
 1
[ CPULongType{7} ]
 0
 4
 7
[ CPULongType{3} ]

But when I load my script model, an error will be thrown.

torch::jit::script::Module model;
std::string file_name = "D:\\TESTProgram\\Libtorch_1.13_cu116\\TorchTest_1.13+cu116\\Dualgnn_module_0409.pt";
try {
    model = torch::jit::load(file_name);
    }
catch (std::exception& e) 
    {
    std::cout << e.what() << std::endl;
    return -1;
    }

Unknown builtin op: torch_scatter::segment_sum_csr.
Could not find any similar ops to torch_scatter::segment_sum_csr. This op may not exist or may not be currently supported in TorchScript.
:
  File "code/__torch__/torch_scatter/segment_csr.py", line 35
    indptr: Tensor,
    out: Optional[Tensor]=None) -> Tensor:
  _10 = ops.torch_scatter.segment_sum_csr(src, indptr, out)
        ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
  return _10
def segment_mean_csr(src: Tensor,
'segment_sum_csr' is being compiled since it was called from 'segment_csr'
Serialized   File "code/__torch__/torch_scatter/segment_csr.py", line 5
    out: Optional[Tensor]=None,
    reduce: str="sum") -> Tensor:
  _0 = __torch__.torch_scatter.segment_csr.segment_sum_csr
  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
  _1 = __torch__.torch_scatter.segment_csr.segment_mean_csr
  _2 = __torch__.torch_scatter.segment_csr.segment_min_csr
'segment_csr' is being compiled since it was called from 'segment'
Serialized   File "code/__torch__/torch_geometric/utils/segment.py", line 4
    ptr: Tensor,
    reduce: str="sum") -> Tensor:
  _0 = __torch__.torch_scatter.segment_csr.segment_csr
  ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
  return _0(src, ptr, None, reduce, )
'segment' is being compiled since it was called from 'MeanAggregation.reduce'
Serialized   File "code/__torch__/torch_geometric/nn/aggr/basic.py", line 22
    reduce: str="sum") -> Tensor:
    _1 = __torch__.torch_geometric.nn.aggr.base.expand_left
    _2 = __torch__.torch_geometric.utils.segment.segment
    ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
    _3 = __torch__.torch_geometric.utils.scatter.scatter
    _4 = uninitialized(Tensor)
'MeanAggregation.reduce' is being compiled since it was called from 'MeanAggregation.forward'
  File "E:\Anaconda\envs\Mesh_Net\lib\site-packages\torch_geometric\nn\aggr\basic.py", line 34
                ptr: Optional[Tensor] = None, dim_size: Optional[int] = None,
                dim: int = -2) -> Tensor:
        return self.reduce(x, index, ptr, dim_size, dim, reduce='mean')
                                                                ~~~~~ <--- HERE
Serialized   File "code/__torch__/torch_geometric/nn/aggr/basic.py", line 12
    dim_size: Optional[int]=None,
    dim: int=-2) -> Tensor:
    _0 = (self).reduce(x, index, ptr, dim_size, dim, "mean", )
                                                     ~~~~~~ <--- HERE
    return _0
  def reduce(self: __torch__.torch_geometric.nn.aggr.basic.MeanAggregation,

See same question at #1718 (comment)
But it's not clear where the problem lies.
What other work can I do? Or do you have any suggestions to solve this problem?

Versions

pytorch: 1.13.0
libtorch:1.13.0
cuda: 11.6
pyg: 2.3.1
torch_scatter: 2.1.1
torch_sparse: 0.6.18

[qiuqiouba]

After trying, I am not sure if it can be considered as a solution to this problem, but at least there are no errors and I can load my script model normally.
I added a step in the demo to establish a link between torch and scatter.

TORCH_LIBRARY(torch_scatter_custom_ops, m) {
    m.def("segment_sum_csr", &segment_sum_csr);
}

[rusty1s]

Thanks for investigating! This mechanism is probably introduced as required in later PyTorch versions, and we should consider doing it automatically on torch-scatter side.

[qiuqiouba]

Looking forward to doing better in the future.