| 335 | |
| 336 | template <typename T> |
| 337 | ValueRefList 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 | |
| 364 | template <typename T> |
| 365 | ValueRefList indexing_rule( |
nothing calls this directly
no test coverage detected