JIT script mode does not work with CPP extension

1

Error message:

Traceback (most recent call last):
  File "reproduce.py", line 7, in <module>
    @torch.jit.script
  File "/home/sidney/anaconda3/envs/py27/lib/python2.7/site-packages/torch/jit/__init__.py", line 616, in script
    graph = _jit_script_compile(ast, _rcb)
  File "/home/sidney/anaconda3/envs/py27/lib/python2.7/site-packages/torch/jit/__init__.py", line 597, in _try_compile_weak_script
    entry = _compiled_weak_fns.get(fn)
  File "/home/sidney/anaconda3/envs/py27/lib/python2.7/weakref.py", line 391, in get
    return self.data.get(ref(key),default)
TypeError: cannot create weak reference to 'builtin_function_or_method' object

Files:

// preprocess.cc

#include <pybind11/pybind11.h>
#include <string>
#include <iostream>

int convert_str_to_tensor(std::string str)
{
  std::cout << str << std::endl;
  return 0;
}

PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
  m.def("convert_str_to_tensor", &convert_str_to_tensor, "LSTM preprocess");
}

// lstm.py

import torch
from torch.utils.cpp_extension import load

preprocess = load(name="preprocess", sources=["preprocess.cc"])

@torch.jit.script
def lstm(cells):
    # type: (List[torch.Tensor]) -> torch.Tensor
    preprocess.convert_str_to_tensor("sidney")

    hidden = torch.ones([10, 10])

    for i in range(len(cells)):
        hidden = hidden.mm(cells[i])

    return hidden

print lstm([
    torch.ones([10, 10]),
    torch.ones([10, 10]),
    torch.ones([10, 10])
])
2

You need custom operators for that.
A hint of how they work can be had from the tests:


Best regards

Thomas