#include "megbrain/opr/tensor_gen.h"
#include "megbrain/opr/io.h"
#include "megbrain/test/autocheck.h"
#include "megbrain/test/helper.h"
using namespace mgb;
using namespace opr;
TEST(TestTensorGen, Alloc) {
auto host_x =
std::make_shared<HostTensorND>(CompNode::load("xpu0"), dtype::Int32());
auto graph = ComputingGraph::make();
auto x = opr::Host2DeviceCopy::make_no_value_infer(*graph, host_x),
y = opr::Alloc::make(x, dtype::Float32());
HostTensorND host_y;
auto func = graph->compile({make_callback_copy(y, host_y)});
auto px = host_x->resize({3}).ptr<int>();
px[0] = 2;
px[1] = 3;
px[2] = 5;
func->execute();
ASSERT_EQ(TensorShape({2, 3, 5}), host_y.shape());
}
TEST(TestTensorGen, Linspace) {
auto host_num =
std::make_shared<HostTensorND>(CompNode::load("xpu0"), dtype::Int32());
host_num->resize({1}).ptr<int>()[0] = 30;
using Checker = AutoOprChecker<2, 1>;
for (auto endpoint : {false, true}) {
auto make_graph =
[endpoint, &host_num](
const Checker::SymInpArray& inputs) -> Checker::SymOutArray {
auto num = opr::Host2DeviceCopy::make(
*inputs[0].node()->owner_graph(), host_num)
.rename("num");
return {opr::Linspace::make(
inputs[0].rename("start"), inputs[1].rename("stop"), num,
{endpoint})
.rename("linspace")};
};
auto fwd = [&](Checker::NumOutArray& dest, Checker::NumInpArray inp) {
size_t num = host_num->ptr<int>()[0];
auto ptr = dest[0].resize({num}).ptr<float>();
auto start = *inp[0]->ptr<float>(), stop = *inp[1]->ptr<float>(),
step = (stop - start) / std::max<int>((num - endpoint), 1);
for (size_t i = 0; i < num; ++i)
ptr[i] = start + step * i;
};
Checker::RunOptions opt;
opt.numdiff_eps = 1; std::array<TensorShape, 2> ishp{TensorShape{1}, {1}};
Checker checker(make_graph, fwd);
host_num->ptr<int>()[0] = 30;
checker.run(ishp, opt).run(ishp, opt);
host_num->ptr<int>()[0] = 1;
checker.run(ishp, opt);
}
}
TEST(TestTensorGen, Eye) {
auto graph = ComputingGraph::make();
auto x = opr::Eye::make(
SymbolVar::make_scalar(5, *graph, CompNode::load("xpu0")),
{-1, DTypeEnum::Int32});
HostTensorND host_x;
auto func = graph->compile({make_callback_copy(x, host_x)});
func->execute();
ASSERT_EQ(TensorShape({5, 5}), host_x.shape());
auto ptr = host_x.ptr<int>();
for (int i = 0; i < 5; ++i) {
for (int j = 0; j < 5; ++j)
ASSERT_EQ(*(ptr++), i - j - 1 == 0);
}
}