Method split_tensor

dnn/src/arm_common/lstm/lstm_utils.cpp:49–61

}

TensorNDArray megdnn::arm_common::split_tensor(
        _megdnn_tensor_in tensor, size_t nr_tensor, const TensorLayout& layout) {
    megdnn_assert(
            tensor.layout.span().dist_byte() == nr_tensor * layout.span().dist_byte());

    TensorNDArray tensors;
    auto ptr = tensor.get_ref_ptr();
    for (size_t i = 0; i < nr_tensor; i++) {
        tensors.push_back(TensorND(layout, ptr));
        ptr += layout.span().dist_byte();
    }
    return tensors;
}

namespace megdnn {
namespace arm_common {