Function TEST

src/opr/test/misc.cpp:45–100

}  // namespace

TEST(TestOprMisc, Argmxx) {
    auto run = [](bool is_max, int32_t axis, TensorShape sshape) {
        auto dshape = sshape;
        dshape.shape[axis] = 1;
        using Checker = AutoOprChecker<1, 1>;
        auto make_graph =
                [&](const Checker::SymInpArray& inputs) -> Checker::SymOutArray {
            if (is_max)
                return {opr::Argmax::make(inputs[0], {axis})};
            else
                return {opr::Argmin::make(inputs[0], {axis})};
        };
        auto better_than = [&](float curr, float best) {
            if (is_max)
                return curr > best;
            else
                return curr < best;
        };
        auto fwd = [&](Checker::NumOutArray& out, Checker::NumInpArray inp) {
            out[0].dtype(dtype::Int32()).resize(dshape);
            size_t A, B, C;
            shape_abc(sshape, axis, A, B, C);
            for (size_t a = 0; a < A; ++a)
                for (size_t c = 0; c < C; ++c) {
                    float best_val;
                    size_t best_arg = -1;
                    if (is_max)
                        best_val = std::numeric_limits<float>::lowest();
                    else
                        best_val = std::numeric_limits<float>::max();
                    for (size_t b = 0; b < B; ++b) {
                        float curr_val = inp[0]->ptr<float>()[(a * B + b) * C + c];
                        if (better_than(curr_val, best_val)) {
                            best_val = curr_val;
                            best_arg = b;
                        }
                    }
                    out[0].ptr<int>()[a * C + c] = best_arg;
                }
        };
        Checker{make_graph, fwd}
                .set_input_allow_grad(0, false)
                .set_output_allow_grad(0, false)
                .run({sshape})
                .run({sshape})
                .run({sshape});
    };
    run(true, 0, {5});
    run(true, 1, {2, 3, 4, 5});
    run(true, 2, {2, 3, 4, 5});
    run(true, 3, {2, 3, 4, 5});
    run(false, 0, {3, 4, 5});
    run(false, 1, {2, 3, 4, 5});
    run(false, 2, {2, 3, 4, 5});
    run(false, 3, {2, 3, 4, 5});
}

TEST(TestOprMisc, Argsort) {