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

Function indexing_rule

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

Source from the content-addressed store, hash-verified

363
364template <typename T>
365ValueRefList indexing_rule(
366 const T& op, Span<ValueRef>& inputs, const bool& auto_convert,
367 const FormatTransformation& t) {
368 mgb_assert(inputs.size() >= 1);
369 auto& src = inputs[0].cast(t.value_type());
370 bool is_reduce_ndim = is_reduce_ndim_idx_items(
371 op.items, {&inputs[1], &inputs[inputs.size() - 1]});
372 if (!is_reduce_ndim) {
373 // only support NHWC2NCHW convert, otherwise maintain src's format
374 if (!(auto_convert && src.format() == FT::NHWC)) {
375 return {t.wrap_output(
376 imperative::apply(op, t.unwrap_inputs(inputs))[0], src.format())};
377 }
378 auto nhwc_items = convert_nchw2nhwc_idx_items(op.items);
379 auto outputs = imperative::apply(
380 *T::make(std::move(nhwc_items), op.scope()), t.unwrap_inputs(inputs));
381 return t.wrap_outputs(outputs, FT::NHWC);
382 }
383 return t.wrap_outputs(imperative::apply(op, t.unwrap_inputs(inputs)));
384}
385
386template <typename T>
387ValueRefList setsubtensor_rule(

Callers

nothing calls this directly

Calls 11

is_reduce_ndim_idx_itemsFunction · 0.85
wrap_outputMethod · 0.80
unwrap_inputsMethod · 0.80
scopeMethod · 0.80
wrap_outputsMethod · 0.80
applyFunction · 0.50
makeFunction · 0.50
sizeMethod · 0.45
castMethod · 0.45
formatMethod · 0.45

Tested by

no test coverage detected