Function broadcast_rule

imperative/src/impl/transformations/format.cpp:256–294

}

ValueRefList broadcast_rule(
        const Broadcast& op, Span<ValueRef>& inputs, const bool& auto_convert,
        const FormatTransformation& t) {
    mgb_assert(inputs.size() >= 1);
    auto& src = inputs[0].cast(t.value_type());
    if (auto_convert && src.format() == FT::NHWC) {
        if (inputs.size() == 1) {
            if (op.shape.size() == 4) {
                // output is still NHWC format
                auto nhwc_shape = convert_nchw2nhwc_vector(op.shape);
                auto outputs = imperative::apply(
                        *Broadcast::make(nhwc_shape), {t.unwrap_input(inputs[0])});
                return t.wrap_outputs(outputs, FT::NHWC);
            } else {
                // will not maintain src's format
                auto nchw_src = t.to(src, FT::DEFAULT, op.scope())->value();
                auto outputs = imperative::apply(op, {nchw_src});
                return t.wrap_outputs(outputs);
            }
        } else if (inputs.size() == 2) {
            auto shape = t.unwrap_input(inputs[1]).numpy()->as_nd();
            if (shape.layout().total_nr_elems() == 4) {
                // output is still NHWC format
                auto nhwc_shape = convert_nchw2nhwc_tensornd(shape);
                auto outputs = imperative::apply(
                        op,
                        SmallVector<ValueRef>{t.unwrap_input(inputs[0]), nhwc_shape});
                return t.wrap_outputs(outputs, FT::NHWC);
            } else {
                // will not maintain src's format
                auto nchw_src = t.to(src, FT::DEFAULT, op.scope())->value();
                auto outputs = imperative::apply(
                        op, SmallVector<ValueRef>{nchw_src, t.unwrap_input(inputs[1])});
                return t.wrap_outputs(outputs);
            }
        }
    }
    return t.wrap_outputs(imperative::apply(op, t.unwrap_inputs(inputs)));
}

inline bool is_reduce_ndim_idx_items(
        const std::vector<std::tuple<int8_t, bool, bool, bool, bool>>& items,