1 #pragma once 2 3 #include <ATen/core/ivalue.h> 4 #include <functional> 5 #include <vector> 6 7 namespace torch::jit::mobile { 8 9 using Stack = std::vector<c10::IValue>; 10 11 void registerPrimOpsFunction( 12 const std::string& name, 13 const std::function<void(Stack&)>& fn); 14 15 bool hasPrimOpsFn(const std::string& name); 16 17 std::function<void(Stack&)>& getPrimOpsFn(const std::string& name); 18 19 class prim_op_fn_register { 20 public: prim_op_fn_register(const std::string & name,const std::function<void (Stack &)> & fn)21 prim_op_fn_register( 22 const std::string& name, 23 const std::function<void(Stack&)>& fn) { 24 registerPrimOpsFunction(name, fn); 25 } 26 }; 27 28 } // namespace torch::jit::mobile 29