test(mgb/index): add empty index desc tests

GitOrigin-RevId: 1a71ad3ede

src/opr/test/indexing.cpp

@@ -336,6 +336,23 @@ TEST(TestOprIndexing, MultiAxisVec) {
     checker.run({TensorShape{12, 1, 2, 9}, {23}}, opt);
 }
 
+TEST(TestOprIndexing, MultiAxisVecWithEmptyIndexDesc) {
+    auto graph = ComputingGraph::make();
+    auto host_x = HostTensorGenerator<>{}({2, 3});
+    auto run_check = [&](SymbolVar y) {
+        HostTensorND host_y;
+        auto func = graph->compile({make_callback_copy(y, host_y)});
+        func->execute();
+        ASSERT_EQ(TensorShape({2, 3}), host_y.shape());
+        func->execute();
+        MGB_ASSERT_TENSOR_EQ(*host_x, host_y);
+    };
+
+    auto x = opr::Host2DeviceCopy::make(*graph, host_x);
+
+    run_check(opr::IndexingMultiAxisVec::make(x, {}));
+}
+
 TEST(TestOprIndexing, IncrMultiAxisVec) {
     using Checker = AutoOprChecker<3, 1>;
     using AIdx = opr::indexing::AxisIndexer;
@@ -429,6 +446,26 @@ TEST(TestOprIndexing, SetMultiAxisVec) {
     checker.run({TensorShape{12, 1, 2, 18}, {1}, {5, 1, 2, 1}}, opt);
 }
 
+TEST(TestOprIndexing, SetMultiAxisVecWithEmptyIndexDesc) {
+    auto graph = ComputingGraph::make();
+    auto host_x = HostTensorGenerator<>{}({2, 3}),
+        host_y = HostTensorGenerator<>{}({2, 3});
+    auto run_check = [&](SymbolVar z) {
+        HostTensorND host_z;
+        auto func = graph->compile({make_callback_copy(z, host_z)});
+        // warning should be printed on the first execution
+        func->execute();
+        ASSERT_EQ(TensorShape({2, 3}), host_z.shape());
+        func->execute();
+        MGB_ASSERT_TENSOR_EQ(host_z, *host_y);
+    };
+
+    auto x = opr::Host2DeviceCopy::make(*graph, host_x),
+         y = opr::Host2DeviceCopy::make(*graph, host_y);
+
+    run_check(opr::IndexingSetMultiAxisVec::make(x, y, {}));
+}
+
 TEST(TestOprIndexing, MultiAxisVecDegenerate) {
     auto graph = ComputingGraph::make();
     auto host_x = HostTensorGenerator<>{}({2, 3}),

src/opr/test/tensor_manip.cpp

@@ -602,6 +602,32 @@ TEST(TestTensorManip, SubtensorNegativeAxis) {
         run({TensorShape{2, 3, 4, 5}});
 }
 
+TEST(TestTensorManip, SubtensorWithEmptyIndexDesc) {
+    using Checker = AutoOprChecker<1, 1>;
+
+    auto make_graph = [&](const Checker::SymInpArray &inputs) ->
+            Checker::SymOutArray {
+        auto x = inputs[0];
+        return {opr::Subtensor::make(x, {})};
+    };
+
+    auto fwd = [](Checker::NumOutArray &dest, Checker::NumInpArray inp) {
+        auto iptr = inp[0]->ptr<float>();
+        auto oshp = inp[0]->shape();
+        auto optr = dest[0].resize(oshp).ptr<float>();
+        for (size_t i = 0, it = oshp.total_nr_elems(); i < it; ++i) {
+            optr[i] = iptr[i];
+        }
+    };
+
+    Checker checker(make_graph, fwd);
+    checker.
+        run({TensorShape{5}}).
+        run({TensorShape{2, 3}}).
+        run({TensorShape{2, 3, 4}}).
+        run({TensorShape{2, 3, 4, 5}});
+}
+
 TEST(TestTensorManip, SubtensorShapeInferForDynAxisIdx) {
     HostTensorGenerator<> gen;
     auto host_x = gen({5, 6, 3});
@@ -1106,6 +1132,24 @@ TEST(TestTensorManip, SetSubtensorDynIdx) {
     MGB_ASSERT_TENSOR_EQ(*host_x, host_y);
 }
 
+TEST(TestTensorManip, SetSubtensorWithEmptyIndexDesc) {
+    HostTensorGenerator<> gen;
+    auto host_x = gen({12}), host_y = gen({12});
+
+    auto graph = ComputingGraph::make();
+    auto x = opr::Host2DeviceCopy::make(*graph, host_x),
+         y = opr::Host2DeviceCopy::make(*graph, host_y),
+         z = opr::SetSubtensor::make(x, y, {});
+
+    ASSERT_TRUE(cg::is_static_var_storage(z.node()));
+    HostTensorND host_z;
+
+    auto func = graph->compile({make_callback_copy(z, host_z)});
+    func->execute();
+
+    MGB_ASSERT_TENSOR_EQ(*host_y, host_z);
+}
+
 TEST(TestTensorManip, IncrSubtensor) {
     using Checker = AutoOprChecker<2, 1>;
     auto make_graph = [](const Checker::SymInpArray &inputs) ->