Caffe2 - C++ API
A deep learning, cross platform ML framework
op.cpp
1 #include <torch/script.h>
2 
3 #include "op.h"
4 
5 #include <cstddef>
6 #include <vector>
7 #include <string>
8 
9 std::vector<torch::Tensor> custom_op(
10  torch::Tensor tensor,
11  double scalar,
12  int64_t repeat) {
13  std::vector<torch::Tensor> output;
14  output.reserve(repeat);
15  for (int64_t i = 0; i < repeat; ++i) {
16  output.push_back(tensor * scalar);
17  }
18  return output;
19 }
20 
21 int64_t custom_op2(std::string s1, std::string s2) {
22  return s1.compare(s2);
23 }
24 
25 static auto registry =
27  // We parse the schema for the user.
28  .op("custom::op", &custom_op)
29  .op("custom::op2", &custom_op2)
30  // User provided schema. Among other things, allows defaulting values,
31  // because we cannot infer default values from the signature. It also
32  // gives arguments meaningful names.
33  .op("custom::op_with_defaults(Tensor tensor, float scalar = 1, int repeat = 1) -> Tensor[]",
34  &custom_op);
RegisterOperators & op(const std::string &name, Implementation &&implementation)
Creates a new operator from a name and implementation function (function pointer or function object/l...
Registration class for new operators.