xref: /aosp_15_r20/external/pytorch/torch/csrc/jit/runtime/custom_operator.h (revision da0073e96a02ea20f0ac840b70461e3646d07c45)
1 #pragma once
2 
3 #include <ATen/core/op_registration/op_registration.h>
4 #include <ATen/core/stack.h>
5 #include <torch/csrc/jit/runtime/operator.h>
6 
7 namespace torch::jit {
8 
9 /// Registration class for new operators. Effectively calls
10 /// `torch::jit::registerOperator` for every supplied operator, but allows doing
11 /// so in the global scope when a `RegisterOperators` object is assigned to a
12 /// static variable.
13 /// Note: This is *not* the custom operator API. If you want to register custom
14 /// operators, take a look at torch::RegisterOperators.
15 struct TORCH_API RegisterOperators {
16   RegisterOperators() = default;
17 
18   /// Registers a vector of already created `Operator`s.
19   /// The operator element is now optional to filter null ops. It's backward
20   /// compatible and works for selective operator registration.
RegisterOperatorsRegisterOperators21   explicit RegisterOperators(std::vector<std::optional<Operator>> operators) {
22     for (std::optional<Operator>& o : operators) {
23       if (o) {
24         registerOperator(std::move(o.value()));
25       }
26     }
27   }
28 };
29 
30 } // namespace torch::jit
31