JIT compile and load custom operations

username: ifgovh
cooked: <p>Hi, I tried to build my custom operation and the compilation worked. When I load my ops, I got the below error. The error even happened when I load the example in documentation. Could you please help me out? Thank you so much!</p>
<p><strong>Load and get error</strong><br>
print(torch.ops.my_ops.SpikeFunction)<br>
Traceback (most recent call last):<br>
File “”, line 1, in<br>
File “/calc/guozhang/anaconda3/lib/python3.8/site-packages/torch/_ops.py”, line 61, in <strong>getattr</strong><br>
op = torch._C._jit_get_operation(qualified_op_name)<br>
RuntimeError: No such operator my_ops::SpikeFunction</p>
<p><strong>The code I used to load the module:</strong></p>
<p>torch.utils.cpp_extension.load(<br>
… name=“SpikeFunction”,<br>
… sources=[“spikefunction.cpp”],<br>
… is_python_module=False,<br>
… verbose=True<br>
… )<br>
Using /home/guozhang/.cache/torch_extensions as PyTorch extensions root…<br>
Emitting ninja build file /home/guozhang/.cache/torch_extensions/SpikeFunction/build.ninja…<br>
Building extension module SpikeFunction…<br>
Allowing ninja to set a default number of workers… (overridable by setting the environment variable MAX_JOBS=N)<br>
[1/2] c++ -MMD -MF spikefunction.o.d -DTORCH_EXTENSION_NAME=SpikeFunction -DTORCH_API_INCLUDE_EXTENSION_H -DPYBIND11_COMPILER_TYPE="_gcc" -DPYBIND11_STDLIB="_libstdcpp" -DPYBIND11_BUILD_ABI="_cxxabi1011" -isystem /calc/guozhang/anaconda3/lib/python3.8/site-packages/torch/include -isystem /calc/guozhang/anaconda3/lib/python3.8/site-packages/torch/include/torch/csrc/api/include -isystem /calc/guozhang/anaconda3/lib/python3.8/site-packages/torch/include/TH -isystem /calc/guozhang/anaconda3/lib/python3.8/site-packages/torch/include/THC -isystem /calc/guozhang/anaconda3/include/python3.8 -D_GLIBCXX_USE_CXX11_ABI=0 -fPIC -std=c++14 -c /home/guozhang/multiplicative_rule/eprop/spikefunction.cpp -o spikefunction.o<br>
[2/2] c++ spikefunction.o -shared -L/calc/guozhang/anaconda3/lib/python3.8/site-packages/torch/lib -lc10 -ltorch_cpu -ltorch -ltorch_python -o SpikeFunction.so<br>
Loading extension module SpikeFunction…</p>
<p><strong>The C++ code I wrote:</strong><br>
<span class="hashtag">#include</span> &lt;torch/all.h&gt;<br>
<span class="hashtag">#include</span> &lt;torch/python.h&gt;</p>
<p>class SpikeFunction : public torch::autograd::Function {<br>
public:<br>
static torch::autograd::variable_list forward(<br>
torch::autograd::AutogradContext* ctx,<br>
torch::autograd::Variable v_scaled,<br>
torch::autograd::Variable dampening_factor) {<br>
// forward calculation<br>
ctx-&gt;save_for_backward({v_scaled,dampening_factor});<br>
return {torch::greater(v_scaled, 0.)};<br>
}</p>
<p>static torch::autograd::variable_list backward(<br>
torch::autograd::AutogradContext* ctx,<br>
torch::autograd::variable_list grad_output) {<br>
// backward calculation<br>
auto saved = ctx-&gt;get_saved_variables();<br>
auto v_scaled = saved[0];<br>
auto dampening_factor = saved[1];<br>
auto dE_dz = grad_output[0];<br>
auto dz_dv_scaled = torch::maximum(1 - torch::abs(v_scaled), torch::zeros_like(v_scaled)) * dampening_factor;<br>
auto dE_dv_scaled = dE_dz * dz_dv_scaled;<br>
return {dE_dv_scaled, torch::zeros_like(dampening_factor)};<br>
}<br>
};</p>
<p>torch::autograd::variable_list SpikeFunction(const torch::Tensor&amp; v_scaled, const torch::Tensor&amp; dampening_factor) {<br>
return SpikeFunction::apply(v_scaled,dampening_factor);<br>
}</p>

username: googlebot
cooked: <p>seems that you’re not registering the operator:</p>
<pre><code class="lang-auto">static auto registry =
torch::RegisterOperators().op("myops::SpikeFunction", &amp;SpikeFunction);
</code></pre>

username: ifgovh
cooked: <p>Thank you for your suggestion. Although I don’t why it did work in my case, I checked the keyword “register” and found I should add the below code at the end. It works.</p>
<blockquote>
<p>TORCH_LIBRARY(my_ops, m) {<br>
m.def(“SpikeFunction”, &amp;SpikeFunction);<br>
}</p>
</blockquote>
<p>But I have a following question; the returned tensor’s requires_grad is false and I cannot backward it properly. Initially, I thought it may be due to the logical format then I tried a square operation but the issue emerged.</p>
<blockquote>
<blockquote>
<blockquote>
<blockquote>
<p>a=torch.rand(5,requires_grad=True)<br>
b=torch.ops.my_ops.SpikeFunction(a,torch.tensor(0.3))<br>
b[0]<br>
tensor([True, True, True, True, True])<br>
b[0].requires_grad<br>
False</p>
</blockquote>
</blockquote>
</blockquote>
</blockquote>
<p>By the way, may I ask how to cast tensor from logic to float in c++? I checked the <a href="https://pytorch.org/cppdocs/api/function_namespaceat_1ab6fc23541b8829007e8a219d1e2e10f2.html" rel="noopener nofollow ugc">doc</a> but I don’t know how to use it. I tried torch::greater(v_scaled, 0.)._cast_Float() ,   torch::greater(v_scaled, 0.) .at::_cast_Float() ,  at ::_cast_Float( torch::greater(v_scaled, 0.)) and many other patterns, all do not work.</p>
<p>Thank you!</p>

username: googlebot
cooked: <p>I think that’s another registration type, you need operator registry for autograd to handle a function, as in my snippet.</p>
<blockquote>
<p>By the way, may I ask how to cast tensor from logic to float in c++?</p>
</blockquote>
<p>python ops are normally duplicated in c++, in this case at::Tensor::to(ScalarType) overload would do that.</p>

username: ifgovh
cooked: <p>Thank you. I did what you suggested but it still does not have grad_fun in the output tensor. Could you please try the below c++ code on your side?</p>
<pre><code class="lang-auto">class SpikeFunction : public torch::autograd::Function&lt;SpikeFunction&gt; {
 public:
  static torch::Tensor forward(
      torch::autograd::AutogradContext* ctx,
      torch::autograd::Variable v_scaled,
      torch::Tensor dampening_factor) {
    // forward calculation
    ctx-&gt;save_for_backward({v_scaled,dampening_factor});
    return (v_scaled &gt; 0).type_as(v_scaled);
  }

  static torch::autograd::variable_list backward(
      torch::autograd::AutogradContext* ctx,
      torch::autograd::variable_list grad_output) {
    // backward calculation
    auto saved = ctx-&gt;get_saved_variables();
    auto v_scaled = saved[0];
    auto dampening_factor = saved[1];
    auto dE_dz = grad_output[0];
    auto dz_dv_scaled = torch::maximum(1 - torch::abs(v_scaled), torch::zeros_like(v_scaled)) * dampening_factor;
    auto dE_dv_scaled = dE_dz * dz_dv_scaled;
    return {dE_dv_scaled, torch::Tensor()};
  }
};

torch::Tensor SpikeFunction(torch::Tensor&amp; v_scaled, torch::Tensor&amp; dampening_factor) {
  return SpikeFunction::apply(v_scaled,dampening_factor);
}

// TORCH_LIBRARY(my_ops, m) {
//   m.def("SpikeFunction", &amp;SpikeFunction);
// }

static auto registry = torch::RegisterOperators().op("myops::SpikeFunction", &amp;SpikeFunction);

</code></pre>

username: googlebot
cooked: <p>worked for me on 1.8 build:</p>
<p>torch.ops.myops.SpikeFunction(torch.ones(1).requires_grad_(), torch.ones(1).requires_grad_())<br>
Out[10]: tensor(cpu,(1,)[1.], grad_fn=&lt;CppNode&gt;)</p>
<p>Note that at least one of tensors must have requires_grad, otherwise backward pass is omitted and grad_fn is not set.</p>

username: ifgovh
cooked: <p>Got it. It works! Thank you so much! Have a great weekend!</p>