| 32 | } |
| 33 | |
| 34 | ValueRefList bn1d_rule(const OpDef& op, Span<ValueRef> inputs) { |
| 35 | size_t ndim = inputs.at(0).shape()->ndim; |
| 36 | CompNode device = *inputs.at(0).device(); |
| 37 | bool need_expand = (ndim == 2 || ndim == 3); |
| 38 | if (!need_expand) |
| 39 | return imperative::apply(op, inputs); |
| 40 | |
| 41 | ValueRefList converted(inputs.size()); |
| 42 | std::vector<int32_t> axis = {(int32_t)3}; |
| 43 | if (ndim == 2) { |
| 44 | axis.insert(axis.begin(), (int32_t)2); |
| 45 | } |
| 46 | converted[0] = imperative::apply(ApplyOp(*AddAxis::make(axis)), inputs[0])[0]; |
| 47 | if (device.device_type() == CompNode::DeviceType::CAMBRICON) { |
| 48 | converted[0] = imperative::apply(SetFormat("nhwc"), converted[0])[0]; |
| 49 | } |
| 50 | for (size_t i = 1; i < inputs.size(); ++i) { |
| 51 | converted[i] = inputs[i]; |
| 52 | } |
| 53 | std::reverse(std::begin(axis), std::end(axis)); |
| 54 | auto outputs = imperative::apply(op, converted); |
| 55 | size_t idx = outputs.size() - 1; |
| 56 | if (device.device_type() == CompNode::DeviceType::CAMBRICON) { |
| 57 | outputs[idx] = imperative::apply(SetFormat("nchw"), outputs[idx])[0]; |
| 58 | } |
| 59 | outputs[idx] = imperative::apply(ApplyOp(*RemoveAxis::make(axis)), outputs[idx])[0]; |
| 60 | return outputs; |
| 61 | } |
| 62 | |
| 63 | struct DimExpansionRuleRegistry { |
| 64 | DimExpansionRuleRegistry() { |