Function TEST

src/opr/test/tensor_manip.cpp:16–58

using namespace opr;

TEST(TestTensorManip, GetVarShape) {
    HostTensorGenerator<> gen;
    auto host_x = gen({3, 1}), host_y = gen({1, 2});
    auto graph = ComputingGraph::make();
    auto x = opr::Host2DeviceCopy::make(*graph, host_x),
         y = opr::Host2DeviceCopy::make(*graph, host_y),
         z0 = opr::GetVarShape::make({x, y, x.make_scalar(5)}),
         z1 = opr::GetVarShape::make({x, y}, 1);

    // ensure scalar is removed
    ASSERT_EQ(2u, z0.node()->owner_opr()->input().size());

    constexpr auto tdt = cg::OperatorNodeBase::NodeProp::DepType::SHAPE;
    auto&& dt = z0.node()->owner_opr()->node_prop().dep_map();
    ASSERT_EQ(2u, dt.size());
    ASSERT_EQ(tdt, dt.at(x.node()));
    ASSERT_EQ(tdt, dt.at(y.node()));

    auto as_shp = [](const HostTensorND& hv) {
        mgb_assert(hv.dtype() == dtype::Int32());
        mgb_assert(hv.shape().ndim == 1);
        TensorShape ret;
        ret.ndim = hv.shape()[0];
        auto p = hv.ptr<int>();
        for (size_t i = 0; i < ret.ndim; ++i)
            ret[i] = p[i];
        return ret;
    };
    HostTensorND host_z0, host_z1;
    auto func = graph->compile(
            {make_callback_copy(z0, host_z0), make_callback_copy(z1, host_z1)});
    func->execute();

    ASSERT_EQ(TensorShape({3, 2}), as_shp(host_z0));
    ASSERT_EQ(TensorShape({2}), as_shp(host_z1));

    *host_x = *gen({5, 1, 6});
    *host_y = *gen({1, 8, 1});
    func->execute();

    ASSERT_EQ(TensorShape({5, 8, 6}), as_shp(host_z0));
    ASSERT_EQ(TensorShape({8}), as_shp(host_z1));
}

TEST(TestTensorManip, GetVarShapeBypass) {
    HostTensorGenerator<> gen;