Method dimshuffle

dnn/src/common/basic_types.cpp:231–243

}

TensorLayout TensorLayout::dimshuffle(const std::vector<size_t>& dims) const {
    TensorLayout res{dtype, format};
    res.ndim = this->ndim;
    megdnn_assert_eq_size_t(dims.size(), this->ndim);
    auto ndim = this->ndim;
    rep(i, ndim) {
        auto dest = dims[i];
        megdnn_assert(dest < ndim);
        res.shape[i] = this->shape[dest];
        res.stride[i] = this->stride[dest];
    }
    return res;
}

TensorLayout TensorLayout::remove_axis(size_t idx) const {
    TensorLayout res{*this};