MCPcopy Create free account
hub / github.com/MegEngine/MegEngine / reshape_rule

Function reshape_rule

imperative/src/impl/transformations/format.cpp:215–254  ·  view source on GitHub ↗

Source from the content-addressed store, hash-verified

213}
214
215ValueRefList reshape_rule(
216 const Reshape& op, Span<ValueRef>& inputs, const bool& auto_convert,
217 const FormatTransformation& t) {
218 mgb_assert(inputs.size() >= 1);
219 auto& src = inputs[0].cast(t.value_type());
220 if (auto_convert && src.format() == FT::NHWC) {
221 if (inputs.size() == 1) {
222 if (op.shape.size() == 4) {
223 // output is still NHWC format
224 auto nhwc_shape = convert_nchw2nhwc_vector(op.shape);
225 auto outputs = imperative::apply(
226 *Reshape::make(op.axis, nhwc_shape),
227 {t.unwrap_input(inputs[0])});
228 return t.wrap_outputs(outputs, FT::NHWC);
229 } else {
230 // will not maintain src's format
231 auto nchw_src = t.to(src, FT::DEFAULT, op.scope())->value();
232 auto outputs = imperative::apply(op, {nchw_src});
233 return t.wrap_outputs(outputs);
234 }
235 } else if (inputs.size() == 2) {
236 auto shape = t.unwrap_input(inputs[1]).numpy()->as_nd();
237 if (shape.layout().total_nr_elems() == 4) {
238 // output is still NHWC format
239 auto nhwc_shape = convert_nchw2nhwc_tensornd(shape);
240 auto outputs = imperative::apply(
241 op,
242 SmallVector<ValueRef>{t.unwrap_input(inputs[0]), nhwc_shape});
243 return t.wrap_outputs(outputs, FT::NHWC);
244 } else {
245 // will not maintain src's format
246 auto nchw_src = t.to(src, FT::DEFAULT, op.scope())->value();
247 auto outputs = imperative::apply(
248 op, SmallVector<ValueRef>{nchw_src, t.unwrap_input(inputs[1])});
249 return t.wrap_outputs(outputs);
250 }
251 }
252 }
253 return t.wrap_outputs(imperative::apply(op, t.unwrap_inputs(inputs)));
254}
255
256ValueRefList broadcast_rule(
257 const Broadcast& op, Span<ValueRef>& inputs, const bool& auto_convert,

Callers

nothing calls this directly

Calls 15

convert_nchw2nhwc_vectorFunction · 0.85
unwrap_inputMethod · 0.80
wrap_outputsMethod · 0.80
scopeMethod · 0.80
as_ndMethod · 0.80
unwrap_inputsMethod · 0.80
applyFunction · 0.50
makeFunction · 0.50
sizeMethod · 0.45
castMethod · 0.45
formatMethod · 0.45

Tested by

no test coverage detected