pytorch

1
#include <torch/csrc/jit/mobile/prim_ops_registery.h>
2

3
namespace torch {
4
namespace jit {
5
namespace mobile {
6

7
static std::unordered_map<std::string, std::function<void(Stack&)>>&
8
primOpsFnTable() {
9
  static std::unordered_map<std::string, std::function<void(Stack&)>>
10
      prim_ops_fn;
11
  return prim_ops_fn;
12
}
13

14
void registerPrimOpsFunction(
15
    const std::string& name,
16
    const std::function<void(Stack&)>& fn) {
17
  primOpsFnTable()[name] = fn;
18
}
19

20
bool hasPrimOpsFn(const std::string& name) {
21
  return primOpsFnTable().count(name);
22
}
23

24
std::function<void(Stack&)>& getPrimOpsFn(const std::string& name) {
25
  TORCH_CHECK(
26
      hasPrimOpsFn(name),
27
      "Prim Ops Function for ",
28
      name,
29
      " is not promoted yet.");
30
  return primOpsFnTable()[name];
31
}
32

33
} // namespace mobile
34
} // namespace jit
35
} // namespace torch
36