| 10 | namespace general_norm { |
| 11 | |
| 12 | cg::OperatorNodeBase* apply_on_var_node(const OpDef& def, const VarNodeArray& inputs) { |
| 13 | auto&& op = static_cast<const GeneralNorm&>(def); |
| 14 | size_t nr_inp = inputs.size(); |
| 15 | auto p = op.param(); |
| 16 | mgb_assert((nr_inp == 3 && p.affine) || (nr_inp == 1 && !p.affine)); |
| 17 | OperatorNodeConfig config{op.make_name()}; |
| 18 | if (nr_inp == 3) { |
| 19 | return opr::GeneralNorm::make( |
| 20 | inputs[0], inputs[1], inputs[2], op.param(), config)[0] |
| 21 | .node() |
| 22 | ->owner_opr(); |
| 23 | } else { |
| 24 | return opr::GeneralNorm::make(inputs[0], op.param(), config)[0] |
| 25 | .node() |
| 26 | ->owner_opr(); |
| 27 | } |
| 28 | } |
| 29 | |
| 30 | std::tuple<SmallVector<LogicalTensorDesc>, bool> infer_output_attrs_fallible( |
| 31 | const OpDef& def, const SmallVector<LogicalTensorDesc>& inputs) { |