44 lines
1.2 KiB
C++
44 lines
1.2 KiB
C++
// sherpa-onnx/csrc/pad-sequence-test.cc
|
|
//
|
|
// Copyright (c) 2023 Xiaomi Corporation
|
|
|
|
#include "sherpa-onnx/csrc/pad-sequence.h"
|
|
|
|
#include <numeric>
|
|
|
|
#include "gtest/gtest.h"
|
|
#include "sherpa-onnx/csrc/onnx-utils.h"
|
|
|
|
namespace sherpa_onnx {
|
|
|
|
TEST(PadSequence, ThreeTensors) {
|
|
Ort::AllocatorWithDefaultOptions allocator;
|
|
|
|
std::array<int64_t, 2> shape1{3, 5};
|
|
Ort::Value v1 =
|
|
Ort::Value::CreateTensor<float>(allocator, shape1.data(), shape1.size());
|
|
float *p1 = v1.GetTensorMutableData<float>();
|
|
std::iota(p1, p1 + shape1[0] * shape1[1], 0);
|
|
|
|
std::array<int64_t, 2> shape2{4, 5};
|
|
Ort::Value v2 =
|
|
Ort::Value::CreateTensor<float>(allocator, shape2.data(), shape2.size());
|
|
float *p2 = v2.GetTensorMutableData<float>();
|
|
std::iota(p2, p2 + shape2[0] * shape2[1], 0);
|
|
|
|
std::array<int64_t, 2> shape3{2, 5};
|
|
Ort::Value v3 =
|
|
Ort::Value::CreateTensor<float>(allocator, shape3.data(), shape3.size());
|
|
float *p3 = v3.GetTensorMutableData<float>();
|
|
std::iota(p3, p3 + shape3[0] * shape3[1], 0);
|
|
|
|
auto ans = PadSequence(allocator, {&v1, &v2, &v3}, -1);
|
|
|
|
Print2D(&v1);
|
|
Print2D(&v2);
|
|
Print2D(&v3);
|
|
Print3D(&ans);
|
|
}
|
|
|
|
} // namespace sherpa_onnx
|