Function dimshuffle_rule

imperative/src/impl/transformations/format.cpp:182–196

}

ValueRefList dimshuffle_rule(
        const Dimshuffle& op, Span<ValueRef>& inputs, const bool& auto_convert,
        const FormatTransformation& t) {
    mgb_assert(inputs.size() == 1);
    auto& src = inputs[0].cast(t.value_type());
    // Only support converting pattern from NCHW to NHWC currently.
    if (auto_convert && src.format() == FT::NHWC) {
        auto pattern = convert_nchw2nhwc_pattern(op.pattern);
        // dimshuffle will not maintain NHWC Format
        return t.wrap_outputs(imperative::apply(
                *Dimshuffle::make(std::move(pattern), op.scope()),
                t.unwrap_inputs(inputs)));
    }
    return t.wrap_outputs(imperative::apply(op, t.unwrap_inputs(inputs)));
}

ValueRef convert_nchw2nhwc_tensornd(const HostTensorND& shape) {
    mgb_assert(shape.layout().total_nr_elems() == 4);