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

Function subtensor_rule

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

Source from the content-addressed store, hash-verified

335
336template <typename T>
337ValueRefList subtensor_rule(
338 const T& op, Span<ValueRef>& inputs, const bool& auto_convert,
339 const FormatTransformation& t) {
340 mgb_assert(inputs.size() >= 1);
341 auto& src = inputs[0].cast(t.value_type());
342 bool is_reduce_ndim = false;
343 if (inputs.size() > 1) {
344 is_reduce_ndim = is_reduce_ndim_idx_items(
345 op.items, {&inputs[1], &inputs[inputs.size() - 1]});
346 } else {
347 is_reduce_ndim = is_subtensor_reduce_ndim(op.items, op.slice_items);
348 }
349 if (!is_reduce_ndim) {
350 // only support NHWC2NCHW convert, otherwise maintain src's format
351 if (!(auto_convert && src.format() == FT::NHWC)) {
352 return {t.wrap_output(
353 imperative::apply(op, t.unwrap_inputs(inputs))[0], src.format())};
354 }
355 auto nhwc_items = convert_nchw2nhwc_idx_items(op.items);
356 auto outputs = imperative::apply(
357 *T::make(std::move(nhwc_items), op.slice_items, op.scope()),
358 t.unwrap_inputs(inputs));
359 return t.wrap_outputs(outputs, FT::NHWC);
360 }
361 return t.wrap_outputs(imperative::apply(op, t.unwrap_inputs(inputs)));
362}
363
364template <typename T>
365ValueRefList indexing_rule(

Callers

nothing calls this directly

Calls 12

is_reduce_ndim_idx_itemsFunction · 0.85
is_subtensor_reduce_ndimFunction · 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