Function test_softmax_variations

test/op_shape_test.cpp:1610–1645

}

static void test_softmax_variations(const std::string& name)
{
    {
        migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
        expect_shape(migraphx::shape{migraphx::shape::float_type, {2, 3, 4, 5}},
                     migraphx::make_op(name, {{"axis", 0}}),
                     input);
    }

    {
        migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
        expect_shape(migraphx::shape{migraphx::shape::float_type, {2, 3, 4, 5}},
                     migraphx::make_op(name, {{"axis", 1}}),
                     input);
    }

    {
        migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
        expect_shape(migraphx::shape{migraphx::shape::float_type, {2, 3, 4, 5}},
                     migraphx::make_op(name, {{"axis", 2}}),
                     input);
    }

    {
        migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
        expect_shape(migraphx::shape{migraphx::shape::float_type, {2, 3, 4, 5}},
                     migraphx::make_op(name, {{"axis", 3}}),
                     input);
    }

    {
        migraphx::shape input{migraphx::shape::float_type, {2, 3, 4, 5}};
        int axis = 4;
        throws_shape(migraphx::make_op(name, {{"axis", axis}}), input);
    }
}
TEST_CASE(logsoftmax) { test_softmax_variations("logsoftmax"); }

TEST_CASE(softmax) { test_softmax_variations("softmax"); }