How to register our custom operator from a shared library like C++

[FelixHo]

I customized some operators (via C++) and called them in the JIT model.

from torch.utils.cpp_extension import load_inline
import torch
from typing import *

op_source = """
int64_t sum_loop(int64_t n) 
{
    int64_t sum = 0;
    for (int64_t i = 0; i < n; i++){
        sum+=1;
    }
    return sum;
}


TORCH_LIBRARY(my_ops, m) {
  m.def("sum_loop_cpp", &sum_loop);
}
"""
load_inline(name="my_ops", cpp_sources=op_source, is_python_module=False, verbose=True)

@torch.jit.script
def func(x:int) -> int:
    return torch.ops.my_ops.sum_loop_cpp(x)

func.save('/tmp/func.pt')

jit_func = torch.jit.load('/tmp/func.pt')
jit_func(100000) # 100000

when load the JIT model through gotch(ts.ModuleLoad),

package main

import (
	"fmt"

	ts "github.com/sugarme/gotch/tensor"
)

func main() {
    a := []ts.IValue{*ts.NewIValue(int64(100000))}
    func, _ := ts.ModuleLoad("/tmp/func.pt")
    out, err := func.ForwardIs(a)
    if err != nil {
        fmt.Println(err)
    } else {
        fmt.Println(out)
    }
}

I think since I did not register my custom op in runtime, I got an exception

Libtorch API Error:
Unknown builtin op: my_ops::sum_loop_cpp.
Could not find any similar ops to my_ops::sum_loop_cpp. This op may not exist or may not be currently supported in TorchScript.
:
  File "<ipython-input-1-ad0082b9b01e>", line 24
@torch.jit.script
def func(x:int) -> int:
    return torch.ops.my_ops.sum_loop_cpp(x)
           ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE
Serialized   File "code/__torch__.py", line 7
  def forward(self: __torch__.PlaceholderModule,
    x: int) -> int:
    return ops.my_ops.sum_loop_cpp(x)
           ~~~~~~~~~~~~~~~~~~~~~~~ <--- HERE

how can I register the custom ops correctly, similar to the C++ case introduced in this document:
https://pytorch.org/tutorials/advanced/torch_script_custom_ops.html#using-the-torchscript-custom-operator-in-c

[sugarme]

@FelixHo ,

Again, this error message came from libtorch side and unfortunately not with gotch. Golang Cgo might deal with C code in general but it is not our priority at the moment. We might look into it if something changed.

I close for now, feel free to reopen if you have any ideas/solutions. Thanks.