| 96 | } |
| 97 | |
| 98 | cg::OperatorNodeBase* apply_on_var_node(const OpDef& def, const VarNodeArray& inputs) { |
| 99 | auto&& bn_opr = def.cast_final_safe<BatchNorm>(); |
| 100 | size_t nr_inp = inputs.size(); |
| 101 | mgb_assert( |
| 102 | nr_inp == 3 || nr_inp == 5, |
| 103 | "BatchNorm expects 3 or 5 inputs; got %lu actually", nr_inp); |
| 104 | OperatorNodeConfig config{bn_opr.make_name()}; |
| 105 | if (nr_inp == 3) { |
| 106 | return opr::BatchNorm::make( |
| 107 | inputs[0], inputs[1], inputs[2], bn_opr.param(), config)[0] |
| 108 | .node() |
| 109 | ->owner_opr(); |
| 110 | } else { |
| 111 | return opr::BatchNorm::make( |
| 112 | inputs[0], inputs[1], inputs[2], inputs[3], inputs[4], |
| 113 | bn_opr.param(), config)[0] |
| 114 | .node() |
| 115 | ->owner_opr(); |
| 116 | } |
| 117 | } |
| 118 | |
| 119 | std::tuple<SmallVector<LogicalTensorDesc>, bool> infer_output_attrs_fallible( |
| 120 | const OpDef& def, const SmallVector<LogicalTensorDesc>& inputs) { |