Function setsubtensor_rule

imperative/src/impl/transformations/format.cpp:387–424

template <typename T>
ValueRefList setsubtensor_rule(
        const T& op, Span<ValueRef>& inputs, const bool& auto_convert,
        const FormatTransformation& t) {
    mgb_assert(inputs.size() >= 2);
    auto& src = inputs[0].cast(t.value_type());
    bool is_reduce_ndim = is_reduce_ndim_idx_items(
            op.items, {&inputs[2], &inputs[inputs.size() - 1]});
    if (!is_reduce_ndim) {
        // only support NHWC2NCHW convert, otherwise maintain src's format
        if (!(auto_convert && src.format() == FT::NHWC)) {
            return {t.wrap_output(
                    imperative::apply(op, t.unwrap_inputs(inputs))[0], src.format())};
        }
        // value has been broadcasted to src's fake NCHW shape.
        auto& value = inputs[1].cast(t.value_type());
        auto& format = value.format();
        auto nhwc_inputs = ValueRefList(inputs.size());
        if (format == FT::DEFAULT || format == FT::NCHW) {
            // value for setsubtensor should transpose to match shape.
            auto nhwc_value = t.to(value, FT::NHWC);
            // make new inputs for setsubtensor
            nhwc_inputs[0] = src.value();
            nhwc_inputs[1] = nhwc_value->value();
            for (auto i = 2; i < inputs.size(); ++i) {
                nhwc_inputs[i] = t.unwrap_input(inputs[i]);
            }
        } else if (format != FT::NHWC) {
            mgb_throw(
                    MegBrainError, "Unsupported format(%s) of value for setsubtensor.",
                    format.to_string().c_str());
        }
        auto nhwc_items = convert_nchw2nhwc_idx_items(op.items);
        auto outputs = imperative::apply(
                *T::make(std::move(nhwc_items), op.scope()), nhwc_inputs);
        return t.wrap_outputs(outputs, FT::NHWC);
    }
    return t.wrap_outputs(imperative::apply(op, t.unwrap_inputs(inputs)));
}

inline FT get_inputs_format(Span<ValueRef>& inputs, const FormatTransformation& t) {
    FT format(FT::DEFAULT);