forked from jiuyuan/InfiniTensor
25 lines
874 B
C++
25 lines
874 B
C++
|
#include "core/graph.h"
|
||
|
|
#include "core/runtime.h"
|
||
|
|
#include "operators/pad.h"
|
||
|
|
#include "test.h"
|
||
|
|
|
||
|
|
namespace infini {
|
||
|
|
TEST(Pad, ShapeInference) {
|
||
|
|
Runtime cpuRuntime = CpuRuntimeObj::getInstance();
|
||
|
|
{
|
||
|
|
Graph g = make_ref<GraphObj>(cpuRuntime);
|
||
|
|
Tensor i = g->addTensor({1, 64, 162, 162}, DataType::UInt32);
|
||
|
|
auto op = g->addOp<PadObj>(
|
||
|
|
i, nullptr, vector<int>{2, 10, 1, 5, 0, 10, 1, 5}, std::nullopt);
|
||
|
|
EXPECT_EQ(op->getOutput()->getDims(), (Shape{3, 84, 164, 172}));
|
||
|
|
}
|
||
|
|
{
|
||
|
|
Graph g = make_ref<GraphObj>(cpuRuntime);
|
||
|
|
Tensor i = g->addTensor({1, 64, 162, 162}, DataType::UInt32);
|
||
|
|
auto op = g->addOp<PadObj>(i, nullptr, vector<int>{2, 10, 1, 5},
|
||
|
|
vector<int>{0, 3});
|
||
|
|
EXPECT_EQ(op->getOutput()->getDims(), (Shape{4, 64, 162, 177}));
|
||
|
|
}
|
||
|
|
}
|
||
|
|
|
||
|
|
} // namespace infini
|