pad-sequence-test.cc
1.2 KB
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
// 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