|
|
|
|
@ -30,23 +30,23 @@ using std::endl;
|
|
|
|
|
|
|
|
|
|
void CreateInput(LoDTensor* ids, LoDTensor* scores) {
|
|
|
|
|
LoD lod;
|
|
|
|
|
vector<size_t> level0{0, 2, 4};
|
|
|
|
|
vector<size_t> level1{0, 1, 2, 3, 4};
|
|
|
|
|
vector<size_t> level0({0, 2, 4});
|
|
|
|
|
vector<size_t> level1({0, 1, 2, 3, 4});
|
|
|
|
|
lod.push_back(level0);
|
|
|
|
|
lod.push_back(level1);
|
|
|
|
|
ids->set_lod(lod);
|
|
|
|
|
scores->set_lod(lod);
|
|
|
|
|
|
|
|
|
|
auto dims = framework::make_ddim(vector<int64_t>{4, 3});
|
|
|
|
|
auto dims = framework::make_ddim(vector<int64_t>({4, 3}));
|
|
|
|
|
ids->Resize(dims);
|
|
|
|
|
scores->Resize(dims);
|
|
|
|
|
CPUPlace place;
|
|
|
|
|
|
|
|
|
|
auto* ids_data = ids->mutable_data<int64_t>(place);
|
|
|
|
|
auto* scores_data = scores->mutable_data<float>(place);
|
|
|
|
|
vector<int64_t> _ids{4, 2, 5, 2, 1, 3, 3, 5, 2, 8, 2, 1};
|
|
|
|
|
vector<float> _scores{0.5f, 0.3f, 0.2f, 0.6f, 0.3f, 0.1f,
|
|
|
|
|
0.9f, 0.5f, 0.1f, 0.7f, 0.5f, 0.1f};
|
|
|
|
|
vector<int64_t> _ids({4, 2, 5, 2, 1, 3, 3, 5, 2, 8, 2, 1});
|
|
|
|
|
vector<float> _scores({0.5f, 0.3f, 0.2f, 0.6f, 0.3f, 0.1f,
|
|
|
|
|
0.9f, 0.5f, 0.1f, 0.7f, 0.5f, 0.1f});
|
|
|
|
|
|
|
|
|
|
for (int i = 0; i < 12; i++) {
|
|
|
|
|
ids_data[i] = _ids[i];
|
|
|
|
|
@ -79,8 +79,8 @@ TEST(DISABLED_beam_search_op, run) {
|
|
|
|
|
|
|
|
|
|
ASSERT_EQ(sids.lod(), sscores.lod());
|
|
|
|
|
|
|
|
|
|
vector<int> tids{4, 2, 3, 8};
|
|
|
|
|
vector<float> tscores{0.5f, 0.6f, 0.9f, 0.7f};
|
|
|
|
|
vector<int> tids({4, 2, 3, 8});
|
|
|
|
|
vector<float> tscores({0.5f, 0.6f, 0.9f, 0.7f});
|
|
|
|
|
|
|
|
|
|
for (int i = 0; i < 4; i++) {
|
|
|
|
|
ASSERT_EQ(tids[i], sids.data<int64_t>()[i]);
|
|
|
|
|
|